RESUNET ARCHITECTURE :
MY IMPLEMENTATION:
import torch
import torch.nn as nn
import torch.nn.functional as TF
# Residual Block
class ResidualBlock(nn.Module):
def __init__(self , in_channels , out_channels , stride=1):
super().__init__()
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
)
# Shortcut connection (identity mapping)
self.shortcut = nn.Sequential()
if stride != 1 or in_channels != out_channels: # Fixed this condition
self.shortcut = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_channels)
)
def forward(self, x):
identity = self.shortcut(x)
out = self.double_conv(x)
out += identity
return torch.relu(out)
# Convolutional Residual Block
class ConvResidual(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv_res = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
ResidualBlock(out_channels, out_channels),
nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
)
def forward(self, x):
return torch.relu(self.conv_res(x))
# Downsampling
class Down(nn.Module):
def __init__(self):
super().__init__()
self.down = nn.MaxPool2d(kernel_size=2, stride=2)
def forward(self, x):
return self.down(x)
# Upsampling
class Up(nn.Module):
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
self.bilinear = bilinear
if self.bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv_residuals = ConvResidual(in_channels=in_channels, out_channels=out_channels)
else:
self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
self.conv_residuals = ConvResidual(in_channels=out_channels * 2, out_channels=out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
diffY = x2.size(2) - x1.size(2)
diffX = x2.size(3) - x1.size(3)
x1 = TF.pad(x1, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2))
out = torch.cat([x2, x1], dim=1)
return self.conv_residuals(out)
# Output Layer
class Out(nn.Module):
def __init__(self, in_channels, n_classes):
super().__init__()
self.out = nn.Sequential(
nn.Conv2d(in_channels, n_classes, kernel_size=1),
nn.Sigmoid()
)
def forward(self, X):
return self.out(X)
# Full ResUNet Model
class ResUNet(nn.Module):
def __init__(self, in_channels, n_channels, n_classes, bilinear=True):
super().__init__()
self.encoder1 = ConvResidual(in_channels, n_channels)
self.down1 = Down()
self.encoder2 = ConvResidual(n_channels, 2 * n_channels)
self.down2 = Down()
self.bottleneck = ConvResidual(2 * n_channels, 4 * n_channels)
self.up1 = Up(4 * n_channels, 2 * n_channels, bilinear=bilinear)
self.decoder1 = ConvResidual(2 * n_channels, n_channels)
self.up2 = Up(2 * n_channels, n_channels, bilinear=bilinear)
self.out = Out(n_channels, n_classes)
def forward(self, x):
x1 = self.encoder1(x)
x1_down = self.down1(x1)
x2 = self.encoder2(x1_down)
x2_down = self.down2(x2)
x_bottleneck = self.bottleneck(x2_down)
x = self.up1(x_bottleneck, x2)
x = self.decoder1(x)
x = self.up2(x, x1)
return self.out(x)
if __name__ == "__main__":
model = ResUNet(3, 64, 3)
x = torch.rand(2, 3, 256, 256) # Fixed input channels (was 128, changed to 3)
y = model(x)
print(f"Output shape: {y.shape}") # Should be (2, 3, 256, 256)
I don’t know how to fix this problem .The output suggests,
RuntimeError: Given groups=1, weight of size [128, 256, 3, 3], expected input[2, 384, 128, 128] to have 256 channels, but got 384 channels instead