@Oli I’m trying to implement a model similar to what I posted here before with some differences. The model should be same as this attached pic. Please note that the kernel size value corresponds to the number of feature maps (dimensionality of the output of the layer).

Capture.PNG580×963 101 KB

here is the code:

class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride):
        super(BaseConv, self).__init__()
        self.act = nn.ReLU()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding, stride)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding, stride)
     
                
    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))

        return x
    
class DownConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride):
        super(DownConv, self).__init__()
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.conv_block = BaseConv(in_channels, out_channels, kernel_size, padding, stride)
    
    def forward(self, x):
        x = self.conv_block(x)
        x = self.pool1(x)
        return x

    
    
class UpConv(nn.Module):
    def __init__ (self, in_channels, in_channels_skip, out_channels, kernel_size, padding, stride):
        super(UpConv, self).__init__()
        self.conv_trans1 = nn.Upsample(size=(in_channels, in_channels), scale_factor=(2))
        self.conv_block = BaseConv(in_channels=in_channels + in_channels_skip, out_channels= out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        
    def forward(self, x, x_skip):
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv_trans1(x)
        x= self.conv_block(x)
        return x


class UNet(nn.Module):
    def __init__(self, in_channels, out_channels, n_class, kernel_size, padding, stride):
        super(UNet, self).__init__()

        self.down1 = DownConv(in_channels, out_channels, kernel_size, padding, stride)

        self.down2 = DownConv(out_channels, out_channels, kernel_size, padding, stride)

        self.down3 = DownConv(out_channels, 2 * out_channels, kernel_size, padding, stride)

        self.down4 = DownConv(2 * out_channels, 4 * out_channels, kernel_size, padding, stride)
        
        self.down5 = DownConv(4 * out_channels, 4 * out_channels, kernel_size, padding, stride)
        
        
        self.down6 = BaseConv( 4 * out_channels, 4 * out_channels, kernel_size, padding, stride)
        

        self.up1 = UpConv(4 * out_channels, 4 * out_channels, 4 * out_channels, kernel_size, padding, stride)

        self.up2 = UpConv(4 * out_channels, 4 * out_channels, 4 * out_channels, kernel_size, padding, stride)
        
        self.up3 = UpConv(4 * out_channels, 2 * out_channels, 2 * out_channels, kernel_size, padding, stride)

        self.up4 = UpConv(2 * out_channels, out_channels, out_channels, kernel_size, padding, stride)

        self.up5 = UpConv(out_channels, out_channels, out_channels, kernel_size, padding, stride)
        
      
        self.out = nn.Conv2d(out_channels, n_class, kernel_size, padding, stride)


    def forward(self, x):
        # Encoder
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)
        
        x6 = self.down6(x5)

        # Decoder
        x_up1 = self.up1(x6, x5)
        x_up2 = self.up2(x_up1, x4)
        x_up3 = self.up3(x_up2, x3)
        x_up4 = self.up4(x_up3, x2)
        x_up5 = self.up5(x_up4, x1)
        
        x_out = F.log_softmax(self.out(x_up5), 1)
        return x_out
 
    
model = UNet(in_channels=1,
             out_channels=32,
             n_class=2,
             kernel_size=3,
             padding=1,
             stride=1)

model = model.to(device)


#Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
     print(param_tensor, "\t", model.state_dict()[param_tensor].size())


I think down layers are Ok, but the problem is up5 which is not correspond with the pic as it doesn’t have 32 output. Could you please point me to the right direction?
Here is the print model tate_dict
Model's state_dict:
down1.conv_block.conv1.weight    torch.Size([32, 1, 3, 3])
down1.conv_block.conv1.bias      torch.Size([32])
down1.conv_block.conv2.weight    torch.Size([32, 32, 3, 3])
down1.conv_block.conv2.bias      torch.Size([32])
down2.conv_block.conv1.weight    torch.Size([32, 32, 3, 3])
down2.conv_block.conv1.bias      torch.Size([32])
down2.conv_block.conv2.weight    torch.Size([32, 32, 3, 3])
down2.conv_block.conv2.bias      torch.Size([32])
down3.conv_block.conv1.weight    torch.Size([64, 32, 3, 3])
down3.conv_block.conv1.bias      torch.Size([64])
down3.conv_block.conv2.weight    torch.Size([64, 64, 3, 3])
down3.conv_block.conv2.bias      torch.Size([64])
down4.conv_block.conv1.weight    torch.Size([128, 64, 3, 3])
down4.conv_block.conv1.bias      torch.Size([128])
down4.conv_block.conv2.weight    torch.Size([128, 128, 3, 3])
down4.conv_block.conv2.bias      torch.Size([128])
down5.conv_block.conv1.weight    torch.Size([128, 128, 3, 3])
down5.conv_block.conv1.bias      torch.Size([128])
down5.conv_block.conv2.weight    torch.Size([128, 128, 3, 3])
down5.conv_block.conv2.bias      torch.Size([128])
down6.conv1.weight       torch.Size([128, 128, 3, 3])
down6.conv1.bias         torch.Size([128])
down6.conv2.weight       torch.Size([128, 128, 3, 3])
down6.conv2.bias         torch.Size([128])
up1.conv_block.conv1.weight      torch.Size([128, 256, 3, 3])
up1.conv_block.conv1.bias        torch.Size([128])
up1.conv_block.conv2.weight      torch.Size([128, 128, 3, 3])
up1.conv_block.conv2.bias        torch.Size([128])
up2.conv_block.conv1.weight      torch.Size([128, 256, 3, 3])
up2.conv_block.conv1.bias        torch.Size([128])
up2.conv_block.conv2.weight      torch.Size([128, 128, 3, 3])
up2.conv_block.conv2.bias        torch.Size([128])
up3.conv_block.conv1.weight      torch.Size([64, 192, 3, 3])
up3.conv_block.conv1.bias        torch.Size([64])
up3.conv_block.conv2.weight      torch.Size([64, 64, 3, 3])
up3.conv_block.conv2.bias        torch.Size([64])
up4.conv_block.conv1.weight      torch.Size([32, 96, 3, 3])
up4.conv_block.conv1.bias        torch.Size([32])
up4.conv_block.conv2.weight      torch.Size([32, 32, 3, 3])
up4.conv_block.conv2.bias        torch.Size([32])
up5.conv_block.conv1.weight      torch.Size([32, 64, 3, 3])
up5.conv_block.conv1.bias        torch.Size([32])
up5.conv_block.conv2.weight      torch.Size([32, 32, 3, 3])
up5.conv_block.conv2.bias        torch.Size([32])
out.weight       torch.Size([2, 32, 3, 3])
out.bias         torch.Size([2])