Implement a model similar to the UNet

(Neda) #1

I’m trying to implement a model which is similar to the Unet based the attached architecture. (Supplementary materials for:
DeepLearningforSegmentationusinganOpenLarge-ScaleDatasetin2DEchocardiography)


I used this implementation but I did change it based on the attached pic.

here is the Unet 2 implementation:

class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride, droup_rate = False):
        super(BaseConv, self).__init__()
        self.act = nn.ReLU()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding, stride)
        self.b1 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding, stride)
        self.b2 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
                
    def forward(self, x):
        x = self.conv1(x)
        x = self.act(self.b1(x))
        x = self.conv2(x)
        x = self.act(self.b2(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.pool1(x)
        x = self.conv_block(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.act = nn.ReLU()
        self.conv_trans1 = nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, padding=0, stride=2)
        self.b3 = nn.BatchNorm2d(in_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv3 = nn.Conv2d(in_channels=in_channels + in_channels_skip, out_channels= out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b4 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv4 = nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b5 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)


    def forward(self, x, x_skip):
        x = self.conv_trans1(x)
        x = self.act(self.b3(x))
        x = self.conv3(x)
        x = torch.cat((x, x_skip), dim=1)
        x = self.act(self.b4(x))
        x = self.conv4(x)
        x = self.act(self.b5(x))
       
        return x
    
     
        
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels, n_class, kernel_size, padding, stride, droup_rate = False):
        super(UNet, self).__init__()

        self.down1 = DownConv(in_channels, out_channels, kernel_size, padding, stride)#48
        
        self.down2 = DownConv(out_channels, 2 * out_channels, kernel_size, padding, stride)#96

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

        self.down4 = DownConv(4 * out_channels, 8 * out_channels, kernel_size, padding, stride)#384
        
        self.down5 = BaseConv(8 * out_channels, 16 * out_channels, kernel_size, padding, stride)#768
               

        
        self.up4 = UpConv(16 * out_channels, 8 * out_channels, 8 * out_channels, kernel_size, padding, stride)
        
        self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,  kernel_size, padding, stride)

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

        self.up1 = UpConv(2 * 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)
        

        # Decoder
        x_up1 = self.up4(x5, x4)
        x_up2 = self.up3(x_up1, x3)
        x_up3 = self.up2(x_up2, x2)
        x_up4 = self.up1(x_up3, x1)

        
        x_out = F.log_softmax(self.out(x_up4), 1)
        print(x_out.size())
        return x_out
 
    
model = UNet(in_channels=1,
             out_channels=48,
             n_class=2,
             kernel_size=3,
             padding=1,
             stride=1)

model = model.to(device)
#print(model)
print("UNet model created")


#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())

Currently, x_up1 causing an error RuntimeError: Given groups=1, weight of size [384, 1152, 3, 3], expected input[1, 768, 30, 40] to have 1152 channels, but got 768 channels instead

I’m not sure where I am doing wrong. any comments would be appreciated.

(Olof Harrysson) #2

Edit: Think I got it wrong the first time. This time I’m concating asap in the UpConv and changed the out_channels in all the UpConv layers so it matches the kernel size from the picture (Kernel / Pool size). Also changed the order in the UpConv forward function. I haven’t doublechecked but let me know if this is what you’re looking for.

import torch
import torch.nn as nn
import torch.nn.functional as F

class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride, droup_rate = False):
        super(BaseConv, self).__init__()
        self.act = nn.ReLU()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding, stride)
        self.b1 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding, stride)
        self.b2 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
                
    def forward(self, x):
        x = self.conv1(x)
        x = self.act(self.b1(x))
        x = self.conv2(x)
        x = self.act(self.b2(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.pool1(x)
        x = self.conv_block(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.act = nn.ReLU()
        self.conv_trans1 = nn.ConvTranspose2d(in_channels + in_channels_skip, out_channels, kernel_size=2, padding=0, stride=2)
        self.b3 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv3 = nn.Conv2d(in_channels=out_channels, out_channels= out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b4 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv4 = nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b5 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)


    def forward(self, x, x_skip):
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv_trans1(x)
        x = self.act(self.b3(x))
        x = self.conv3(x)
        x = self.act(self.b4(x))
        x = self.conv4(x)
        x = self.act(self.b5(x))
        return x
    
     
        
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels, n_class, kernel_size, padding, stride, droup_rate = False):
        super(UNet, self).__init__()

        self.down1 = DownConv(in_channels, out_channels, kernel_size, padding, stride)#48
        
        self.down2 = DownConv(out_channels, 2 * out_channels, kernel_size, padding, stride)#96

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

        self.down4 = DownConv(4 * out_channels, 8 * out_channels, kernel_size, padding, stride)#384
        
        self.down5 = BaseConv(8 * out_channels, 16 * out_channels, kernel_size, padding, stride)#768

        
        self.up4 = UpConv(16 * out_channels, 8 * out_channels, 8 * out_channels, kernel_size, padding, stride)
        
        self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,  kernel_size, padding, stride)

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

        self.up1 = UpConv(2 * 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)

        # Decoder
        x_up1 = self.up4(x5, x4)
        x_up2 = self.up3(x_up1, x3)
        x_up3 = self.up2(x_up2, x2)
        x_up4 = self.up1(x_up3, x1)

        x_out = F.log_softmax(self.out(x_up4), 1)
        return x_out
 
    
model = UNet(in_channels=1,
             out_channels=48,
             n_class=2,
             kernel_size=3,
             padding=1,
             stride=1)

#print(model)
print("UNet model created")


#Print model's state_dict
print("Model's state_dict:")


inp = torch.randn(2, 1, 128, 128)
out = model(inp)
print('input size: {}'.format(inp.size()))
print('output size: {}'.format(out.size()))

(Neda) #3

The model is based on this paper. They provided the supplementary paper which include the above pic I did share. I couldn’t find the link of the supplementary paper to share with you. Shall I share it through Dropbox or one-drive link ?

(Olof Harrysson) #4

Ah, I updated my answer since then. I think I got it right now. Could you confirm?

1 Like
(Neda) #5

@Oli thanks a lot. I did get print of updated model as follows:


down1.conv_block.conv1.weight          torch.Size([48, 1, 3, 3])
down1.conv_block.conv2.weight          torch.Size([48, 48, 3, 3])

down2.conv_block.conv1.weight          torch.Size([96, 48, 3, 3])
down2.conv_block.conv2.weight          torch.Size([96, 96, 3, 3])

down3.conv_block.conv1.weight          torch.Size([192, 96, 3, 3])
down3.conv_block.conv2.weight          torch.Size([192, 192, 3, 3])

down4.conv_block.conv1.weight          torch.Size([384, 192, 3, 3])
down4.conv_block.conv2.weight          torch.Size([384, 384, 3, 3])

down5.conv1.weight          torch.Size([768, 384, 3, 3])
down5.conv2.weight          torch.Size([768, 768, 3, 3])

up4.conv_trans1.weight          torch.Size([1152, 384, 2, 2])
up4.conv3.weight          torch.Size([384, 384, 3, 3])
up4.conv4.weight          torch.Size([384, 384, 3, 3])

up3.conv_trans1.weight          torch.Size([576, 192, 2, 2])
up3.conv3.weight          torch.Size([192, 192, 3, 3])
up3.conv4.weight          torch.Size([192, 192, 3, 3])

up2.conv_trans1.weight          torch.Size([288, 96, 2, 2])
up2.conv3.weight          torch.Size([96, 96, 3, 3])
up2.conv4.weight          torch.Size([96, 96, 3, 3])

up1.conv_trans1.weight          torch.Size([144, 48, 2, 2])
up1.conv3.weight          torch.Size([48, 48, 3, 3])
up1.conv4.weight          torch.Size([48, 48, 3, 3])

out.weight          torch.Size([2, 48, 3, 3])

apparently the number of channel from up4 to up1 isn’t same as pic.

(Olof Harrysson) #6

Np. Aren’t they? Number of filters in U1=384, U2=192, U3=96, U4=48. Thats the same as in the printout you just posted just the variables are named in reverse (up1 = U4).

Edit: Ah, the number of channels isn’t specified in the pic. It’s the # filters (kernel in pic) that is specified. I think… Long time since I learned how conv2d works :stuck_out_tongue:

(Neda) #7

@Oli yeah you are right that up1=up4. In the pic number of channel also specified. if you look at the print for up4.conv_trans1.weight torch.Size([1152, 384, 2, 2]) is 1152 channel but in the pic up1 is 384. I think the concatenation shouldn’t be in the self.conv_trans1 . Also in the pic connection specified in front of conv col. is that right?

(Olof Harrysson) #8

I think that the pic doesn’t specify input channels but rather output channels, meaning the number of filters. Then the printout of up4.conv_trans1.weight torch.Size([1152, 384, 2, 2]) actually does match the picture since the second number (output channels) matches the pic up1= 384.

I don’t understand what you mean by this

But it is very likely that the code I provided is doing something in the wrong order :blush:

(Neda) #9

@Oli if those filter kernels in pic are number of input channels, then how should I change up4 ?

(Olof Harrysson) #10

I don’t know. In that case, maybe someone else can help out?

But lets think about the actual input to the network. The first layer is the D1. If that picture would specify the number of input channels somewhere it would have to write 1 for grayscale and 3 for RGB. It only specifies 48, which has to be output channels. So again, I believe the pic is showing output channels.

1 Like
(Olof Harrysson) #11

Ok so this implementation takes the output in the downsample layers with the asterisks *** and concats that with input of the upsample layers with matching ***. Note that I changed to order of the DownConv forward pass to first do the conv, then the maxpooling

import torch
import torch.nn as nn
import torch.nn.functional as F

class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride, droup_rate = False):
        super(BaseConv, self).__init__()
        self.act = nn.ReLU()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding, stride)
        self.b1 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding, stride)
        self.b2 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
                
    def forward(self, x):
        x = self.conv1(x)
        x = self.act(self.b1(x))
        x = self.conv2(x)
        x = self.act(self.b2(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)
        down = self.pool1(x)
        return x, down


class UpConv(nn.Module):
    def __init__ (self, in_channels, in_channels_skip, out_channels, kernel_size, padding, stride):
        super(UpConv, self).__init__()
        self.act = nn.ReLU()
        self.conv_trans1 = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, padding=0, stride=2)
        self.b3 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv3 = nn.Conv2d(in_channels=in_channels, out_channels= out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b4 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)
        self.conv4 = nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, padding=padding, stride=stride)
        self.b5 = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True)


    def forward(self, x, x_skip):
        x = self.conv_trans1(x)
        x = self.act(self.b3(x))
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv3(x)
        x = self.act(self.b4(x))
        x = self.conv4(x)
        x = self.act(self.b5(x))
        return x
    
     
        
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels, n_class, kernel_size, padding, stride, droup_rate = False):
        super(UNet, self).__init__()

        self.down1 = DownConv(in_channels, out_channels, kernel_size, padding, stride)#48
        
        self.down2 = DownConv(out_channels, 2 * out_channels, kernel_size, padding, stride)#96

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

        self.down4 = DownConv(4 * out_channels, 8 * out_channels, kernel_size, padding, stride)#384
        
        self.down5 = BaseConv(8 * out_channels, 16 * out_channels, kernel_size, padding, stride)#768

        
        self.up4 = UpConv(16 * out_channels, 8 * out_channels, 8 * out_channels, kernel_size, padding, stride)
        
        self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,  kernel_size, padding, stride)

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

        self.up1 = UpConv(2 * 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, down = self.down1(x)
        x2, down = self.down2(down)
        x3, down = self.down3(down)
        x4, down = self.down4(down)
        x5 = self.down5(down)

        # Decoder
        x_up1 = self.up4(x5, x4)
        x_up2 = self.up3(x_up1, x3)
        x_up3 = self.up2(x_up2, x2)
        x_up4 = self.up1(x_up3, x1)

        x_out = F.log_softmax(self.out(x_up4), 1)
        return x_out
 
    
model = UNet(in_channels=1,
             out_channels=48,
             n_class=2,
             kernel_size=3,
             padding=1,
             stride=1)

#print(model)
print("UNet model created")


#Print model's state_dict
print("Model's state_dict:")


inp = torch.randn(2, 1, 128, 128)
out = model(inp)
print('input size: {}'.format(inp.size()))
print('output size: {}'.format(out.size()))


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

@Oli thank you very much for your help. Yes, you are right. the pic shows the output channel. I read in the paper “Please note that the kernel size value corresponds to the number of feature maps (dimensionality of the output of the layer)”.

the snippet you posted first if we want to keep it as pic, the maxpooling should be after conv. also if we agree that those filter size are output channel, still up-sampling sections up4.conv_trans1.weight torch.Size([1152, 384, 2, 2]) giving the wrong number of channels.

in the second model you posted is causing an error RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same which is about x1, down = self.down1(x)

1 Like
(Olof Harrysson) #13

Hey, I’ll just reply quickly now but the second posting I did is running for me. Perhaps you mixed some of my code with your own? The error message you are getting is because the model and input are not on the same device, meaning they arent both on either the CPU or the GPU. Try adding the model = model.to(device) as in your original code

(Neda) #14

@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).

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])
(Olof Harrysson) #15

Does the code you posted run for you? It doesn’t work for me. It never gets to up5. The up1 concatenates the input so the number of channels are 256, whilst the number of input channels in the first upsample is 128. Perhaps thats something?

Btw, what pytorch version are you running?

(Neda) #16

@Oli Yes it does run for me.

(Olof Harrysson) #17

What pytorch version? And have you put any input through the network?

1 Like
(Neda) #18

@Oli it is Torch version: 0.4.1 . Yes, you are right, when give it input it will stop at x_up1 = self.up1(x6, x5) the error is about upsample

return torch._C._nn.upsample_nearest2d(input, _output_size(2)) ValueError: only one of size or scale_factor should be defined

(Olof Harrysson) #19

Try replacing these functions in the code below. I didn’t get the x_out = F.log_softmax(self.out(x_up5), 1) to work but I bet you got that one :wink:

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(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, int(out_channels/2), kernel_size, padding, stride)
        
        self.out = nn.Conv2d(int(out_channels/2), n_class, kernel_size, padding, stride)
1 Like
(Neda) #20

@Oli thank you for this. when I print the model it seems fine in terms of output, but when give it an input It throwing an error
x = torch.cat((x, x_skip), dim=1) RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 1. Got 7 and 14 in dimension 2 at c:\programdata\miniconda3\conda-bld\pytorch_1533100020551\work\aten\src\thc\generic/THCTensorMath.cu:87