Unet+CityScape underfitting(-ish)

deeplearner6: <p>Hi!</p>
<p>I’m having a problem with my Unet implementation, it seems that the network is unable to overfit on the training data (<strong>the CE loss plateaus at ~3.1 and the per pixel accuracy at 40-50%</strong>).<br>
When i view a training image after every epoch on the training data i can see that the model just pretty much does 50% sky 50% road, which i guess hits the local minima for the <strong>Cross Entropy loss</strong>.<br>
The model trains to 40-50% per pixel acc at the end of epoch 1 and then seems to stay there.</p>
<p>I am using Pytorch default <strong>Adam optimizer</strong>, i tried using SGD with lr=0.01/0.005/0.001 but the same problem persisted. The model uses 34 classes (CityScape Finely annotated 5000 images).</p>
<p>Here is my Unet implementation</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.down = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.down(x)
        return x


class DownScale(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(DownScale, self).__init__()
        self.down_block = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        x = self.down_block(x)
        return x



class UpScale(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(UpScale, self).__init__()

        self.upsample = nn.ConvTranspose2d(in_channels, out_channels, 2, stride=2) 
        self.doubleconv = DoubleConv(in_channels, out_channels)
    def forward(self, x1, x2):

        x1 = self.upsample(x1, output_size=x2.shape)
        assert x1.shape == x2.shape, f"up size:{x1.shape}, previous:{x2.shape}"
        x = torch.cat([x2, x1], dim=1)
        x = self.doubleconv(x)
        return x



class Unet(nn.Module):

    def __init__(self, input_channels, num_classes):
        super(Unet, self).__init__()
        self.inconv = DoubleConv(input_channels, 64)
        self.down_block1 = DownScale(64, 128)
        self.down_block2 = DownScale(128, 256)
        self.down_block3 = DownScale(256, 512)
        self.down_block4 = DownScale(512, 1024)
        self.up_block1 = UpScale(1024, 512)
        self.up_block2 = UpScale(512, 256)
        self.up_block3 = UpScale(256, 128)
        self.up_block4 = UpScale(128, 64)
        self.outputmap = nn.Conv2d(64, num_classes, kernel_size=1)


    def forward(self, x):
        x1 = self.inconv(x)
        x2 = self.down_block1(x1)
        x3 = self.down_block2(x2)
        x4 = self.down_block3(x3)
        x5 = self.down_block4(x4)
        x = self.up_block1(x5, x4)
        x = self.up_block2(x, x3)
        x = self.up_block3(x, x2)
        x = self.up_block4(x, x1)
        x = self.outputmap(x)
        y = nn.Softmax(dim=1)(x)
        return y
</code></pre>
<p>Here is an image of the model after the 2nd epoch (it stays pretty much like this even after 3-4 epochs):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/b/b6561247a93ff134fe663edf99469022d799badf.png" data-download-href="https://discuss.pytorch.org/uploads/default/b6561247a93ff134fe663edf99469022d799badf" title="Capture"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b6561247a93ff134fe663edf99469022d799badf_2_690x280.png" alt="Capture" data-base62-sha1="q11j1b925nQhZh66xXwLFOk8LUj" width="690" height="280" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b6561247a93ff134fe663edf99469022d799badf_2_690x280.png, https://discuss.pytorch.org/uploads/default/optimized/2X/b/b6561247a93ff134fe663edf99469022d799badf_2_1035x420.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/b/b6561247a93ff134fe663edf99469022d799badf.png 2x" data-dominant-color="D2D3CA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Capture</span><span class="informations">1146×466 179 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However for a different semantic segmentation dataset, the model seemed to learn:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/1/102e2e61cd2144a89d93635f0e94676942e8288a.png" data-download-href="https://discuss.pytorch.org/uploads/default/102e2e61cd2144a89d93635f0e94676942e8288a" title="61478081_1848210178612388_3146660222651596800_n"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/1/102e2e61cd2144a89d93635f0e94676942e8288a_2_690x207.png" alt="61478081_1848210178612388_3146660222651596800_n" data-base62-sha1="2j8A6a9GHpRp60qmqzJtLeMpuqe" width="690" height="207" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/1/102e2e61cd2144a89d93635f0e94676942e8288a_2_690x207.png, https://discuss.pytorch.org/uploads/default/original/2X/1/102e2e61cd2144a89d93635f0e94676942e8288a.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/1/102e2e61cd2144a89d93635f0e94676942e8288a.png 2x" data-dominant-color="B1B0AE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">61478081_1848210178612388_3146660222651596800_n</span><span class="informations">910×274 138 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>What i have tried:</p>
<ul>
<li>Cross entropy with <strong>weights per class</strong> (0.1-0.5 for sky/road/building class (the most frequent ones), 1.0 for others), but that didn’t really work.</li>
</ul>
<p>To me this seems to be a loss function problem which doesn’t get out of the local minima, but i thought that cross entropy should sort of work. I haven’t tried dice loss or any other losses, if you could point out any tips it would be wonderful.</p>
<p>Thanks!</p>