UNet implementation

GPistre: I’m still in the process of learning, so I’m not sure my implementation is right. Right now it seems the loss becomes nan quickly, while the network output “pixels” become 0 or 1 seemingly randomly. I’m not sure it is because of my implementation or because of my lack of understanding of the loss (I pass the last layer through a LogSoftmax, and then use NLLLoss2d). <SCODE>class UNetConvBlock(nn.Module):
    def __init__(self, in_size, out_size, kernel_size=3, activation=F.relu):
        super(UNetConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_size, out_size, kernel_size)
        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size)
        self.activation = activation

    def forward(self, x):
        out = self.activation(self.conv(x))
        out = self.activation(self.conv2(out))

        return out


class UNetUpBlock(nn.Module):
    def __init__(self, in_size, out_size, kernel_size=3, activation=F.relu, space_dropout=False):
        super(UNetUpBlock, self).__init__()
        self.up = nn.ConvTranspose2d(in_size, out_size, 2, stride=2)
        self.conv = nn.Conv2d(in_size, out_size, kernel_size)
        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size)
        self.activation = activation

    def center_crop(self, layer, target_size):
        batch_size, n_channels, layer_width, layer_height = layer.size()
        xy1 = (layer_width - target_size) // 2
        return layer[:, :, xy1:(xy1 + target_size), xy1:(xy1 + target_size)]

    def forward(self, x, bridge):
        up = self.up(x)
        crop1 = self.center_crop(bridge, up.size()[2])
        out = torch.cat([up, crop1], 1)
        out = self.activation(self.conv(out))
        out = self.activation(self.conv2(out))

        return out


class UNet(nn.Module):
    def __init__(self, imsize):
        super(UNet, self).__init__()
        self.imsize = imsize

        self.activation = F.relu
        
        self.pool1 = nn.MaxPool2d(2)
        self.pool2 = nn.MaxPool2d(2)
        self.pool3 = nn.MaxPool2d(2)
        self.pool4 = nn.MaxPool2d(2)

        self.conv_block1_64 = UNetConvBlock(1, 64)
        self.conv_block64_128 = UNetConvBlock(64, 128)
        self.conv_block128_256 = UNetConvBlock(128, 256)
        self.conv_block256_512 = UNetConvBlock(256, 512)
        self.conv_block512_1024 = UNetConvBlock(512, 1024)

        self.up_block1024_512 = UNetUpBlock(1024, 512)
        self.up_block512_256 = UNetUpBlock(512, 256)
        self.up_block256_128 = UNetUpBlock(256, 128)
        self.up_block128_64 = UNetUpBlock(128, 64)

        self.last = nn.Conv2d(64, 2, 1)


    def forward(self, x):
        block1 = self.conv_block1_64(x)
        pool1 = self.pool1(block1)

        block2 = self.conv_block64_128(pool1)
        pool2 = self.pool2(block2)

        block3 = self.conv_block128_256(pool2)
        pool3 = self.pool3(block3)

        block4 = self.conv_block256_512(pool3)
        pool4 = self.pool4(block4)

        block5 = self.conv_block512_1024(pool4)

        up1 = self.up_block1024_512(block5, block4)

        up2 = self.up_block512_256(up1, block3)

        up3 = self.up_block256_128(up2, block2)

        up4 = self.up_block128_64(up3, block1)

        return F.log_softmax(self.last(up4))<ECODE>

apaszke: 

GPistre: Indeed, there’s a typo. A simpler fully convolutional network will also eventually “converge” to nan, albeit more slowly (a few epochs). I was using a similar setup with theano with decent results, so I think my data is ok, but I’m not sure how to troubleshoot. I was using the Adam optimizer, but I’ll try with SGD and a small learning rate and go from there if the results are promising.

apaszke: That’s weird, are you sure the networks are the same and that data is preprocessed correctly? It’s important to find the source of NaNs, that should not happen.

GPistre: It appears there was a size mismatch I overlooked between the network output and the target. (8x2x68x68 vs 8x70x70) I think there might be a bug there as it gives an error when the tensors are on cpu, but is silent and outputs something when they are on gpu. Do you want me to open an issue ? That might be where the nans were coming from. I’ll let it train overnight and see. To reproduce the error: <SCODE>import torch
import torch.nn as nn
from torch.autograd import Variable

loss = nn.NLLLoss2d()

target = Variable(torch.Tensor(8, 70, 70).random_(0, 1)).long()
output = Variable(torch.randn(8, 2, 68, 68))

loss(output, target)
<ECODE> <SCODE>loss(output.cuda(), target.cuda())
<ECODE>

apaszke: Yes please! It seems that this loss is lacking some shape checks on the GPU indeed.

GPistre: last_layer <SCODE>Variable containing:
(0 ,0 ,.,.) = 
    3.5879    3.6678    3.8380  ...     3.1548    3.0576    2.9584
    3.4753    3.7363    3.8944  ...     2.9736    3.0051    2.9889
    3.3298    3.3160    3.5382  ...     2.8276    2.8111    2.8584
              ...                ⋱                ...             
    3.2416    3.1960    3.2502  ...    90.0304   98.9006   98.5473
    3.2719    3.2843    3.2724  ...    39.1980   67.7482   73.4172
    3.2535    3.3880    3.3061  ...    13.5371   25.7164   37.9838

(0 ,1 ,.,.) = 
   -3.7768   -3.8683   -4.0629  ...    -3.2815   -3.1703   -3.0568
   -3.6481   -3.9466   -4.1274  ...    -3.0742   -3.1102   -3.0917
   -3.4816   -3.4659   -3.7200  ...    -2.9072   -2.8883   -2.9424
              ...                ⋱                ...             
   -3.3808   -3.3285   -3.3906  ...  -122.8841 -137.5700 -122.2672
   -3.4154   -3.4296   -3.4160  ...   -48.6680  -82.5093  -91.6747
   -3.3943   -3.5482   -3.4545  ...   -17.6337  -33.1317  -46.8320
[torch.cuda.FloatTensor of size 1x2x68x68 (GPU 0)]
<ECODE> F.log_softmax(last_layer) <SCODE>Variable containing:
(0 ,0 ,.,.) = 
   -0.0006   -0.0005   -0.0004  ...    -0.0016   -0.0020   -0.0024
   -0.0008   -0.0005   -0.0003  ...    -0.0024   -0.0022   -0.0023
   -0.0011   -0.0011   -0.0007  ...    -0.0032   -0.0033   -0.0030
              ...                ⋱                ...             
   -0.0013   -0.0015   -0.0013  ...        nan       nan       nan
   -0.0012   -0.0012   -0.0012  ...     0.0000    0.0000    0.0000
   -0.0013   -0.0010   -0.0012  ...     0.0000    0.0000    0.0000

(0 ,1 ,.,.) = 
   -7.3653   -7.5366   -7.9013  ...    -6.4379   -6.2299   -6.0176
   -7.1242   -7.6833   -8.0221  ...    -6.0501   -6.1175   -6.0829
   -6.8125   -6.7830   -7.2590  ...    -5.7381   -5.7027   -5.8038
              ...                ⋱                ...             
   -6.6237   -6.5260   -6.6420  ...       -inf      -inf      -inf
   -6.6886   -6.7152   -6.6896  ...   -87.8659      -inf      -inf
   -6.6491   -6.9371   -6.7618  ...   -31.1708  -58.8481  -84.8159
[torch.cuda.FloatTensor of size 1x2x68x68 (GPU 0)]
<ECODE> F.log_softmax(last_layer.cpu()) <SCODE>Variable containing:
(0 ,0 ,.,.) = 
   -0.0006   -0.0005   -0.0004  ...    -0.0016   -0.0020   -0.0024
   -0.0008   -0.0005   -0.0003  ...    -0.0024   -0.0022   -0.0023
   -0.0011   -0.0011   -0.0007  ...    -0.0032   -0.0033   -0.0030
              ...                ⋱                ...             
   -0.0013   -0.0015   -0.0013  ...     0.0000    0.0000    0.0000
   -0.0012   -0.0012   -0.0012  ...     0.0000    0.0000    0.0000
   -0.0013   -0.0010   -0.0012  ...    -0.0000    0.0000    0.0000

(0 ,1 ,.,.) = 
   -7.3653   -7.5366   -7.9013  ...    -6.4379   -6.2299   -6.0176
   -7.1242   -7.6833   -8.0221  ...    -6.0501   -6.1175   -6.0829
   -6.8125   -6.7830   -7.2590  ...    -5.7381   -5.7027   -5.8038
              ...                ⋱                ...             
   -6.6237   -6.5260   -6.6420  ...  -212.9145 -236.4706 -220.8145
   -6.6886   -6.7152   -6.6896  ...   -87.8659 -150.2575 -165.0919
   -6.6491   -6.9371   -6.7618  ...   -31.1708  -58.8481  -84.8159
[torch.FloatTensor of size 1x2x68x68]<ECODE>

apaszke: 

lopuhin: 

bodokaiser: Do you have any idea what the problem could be? Do you think it maybe because of missing Batch Norm layers?

lopuhin: One thing that is different in my implementation is that I use upsampling instead of transposed convolutions, it worked significantly better in my case. Batch normalization speeds up convergence but is by no means essential, it worked fine without it too. Simpler models also gave okayish results, but UNet was consistently better - in this task the metric was intersection over union, and simple models were giving results in the 0.2-0.3 range (average over 10 classes), while UNet gave 0.4+ without too much tuning.

bodokaiser: 

Saeed_Izadi: 

fmassa: 

Saeed_Izadi: The point is that the problem persists even after removing batchnorm. Additionally,  I’m using batch size of size 16 in my recent efforts . The same model is working under theano/lasagna implementation!

bodokaiser: Hey, I already commented on this issue on GitHub but got some new ideas today which might be worth to check out: Good luck!

fmassa: If you share the lasagne and the pytorch code, I can have a look

Saeed_Izadi: for the 1. I think it is not the problem, because I’m normalizing the groundtuth to the range of {0,1} and the prediction after sigmoid is also in the range [0,1]. The BCELoss also is supposed to work with this range. For 2, yeah maybe there is a problem with ConvTranspose2d. I will replace it with Upsampling and update you.

Saeed_Izadi: 

fmassa: