Loss value on U-net

username: Jason_Yuen
cooked: <p>I am currently working on Semantic segmentation on VOC2012 dataset with U-net. I encounter a problem that loss is dropping a bit slow and not moving after few epoch.<br>
Here is my training function and model</p>
<pre><code class="lang-auto">def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    running_loss = 0.0
    total_train = 0
    running_correct = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        criterion = nn.CrossEntropyLoss()
        loss = criterion(output, target[:, 0])
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        predicted = torch.argmax(output.data, 1)
        total_train += target.nelement()
        running_correct += predicted.eq(target.data).sum().item()
    print(epoch, running_correct / total_train)
    print(running_loss / len(train_loader.dataset))
</code></pre>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class UNetEnc(nn.Module):
    def __init__(self, in_channels, features, out_channels):
        super().__init__()

        self.up = nn.Sequential(
            nn.Conv2d(in_channels, features, 3), #kernal size: 3
            nn.LeakyReLU(True),
            nn.Conv2d(features, features, 3),
            nn.LeakyReLU(True),
            nn.ConvTranspose2d(features, out_channels, 2, stride=2),
            nn.LeakyReLU(True),
        )

    def forward(self, x):
        return self.up(x)

class UNetDec(nn.Module):
    def __init__(self, in_channels, out_channels, dropout=False):
        super().__init__()

        layers = [
            nn.Conv2d(in_channels, out_channels, 3),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3),
            nn.LeakyReLU(inplace=True),
        ]
        if dropout:
            layers += [nn.Dropout(.5)]
        layers += [nn.MaxPool2d(2, stride=2, ceil_mode=True)]

        self.down = nn.Sequential(*layers)

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

class UNet(nn.Module):

    def __init__(self, num_classes):
        super().__init__()

        self.dec1 = UNetDec(3, 64)
        self.dec2 = UNetDec(64, 128)
        self.dec3 = UNetDec(128, 256)
        self.dec4 = UNetDec(256, 512, dropout=True)
        self.center = nn.Sequential(
            nn.Conv2d(512, 1024, 3),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(1024, 1024, 3),
            nn.LeakyReLU(inplace=True),
            nn.Dropout(),
            nn.ConvTranspose2d(1024, 512, 2, stride=2),
            nn.LeakyReLU(inplace=True),
        )
        self.enc4 = UNetEnc(1024, 512, 256)
        self.enc3 = UNetEnc(512, 256, 128)
        self.enc2 = UNetEnc(256, 128, 64)
        self.enc1 = nn.Sequential(
            nn.Conv2d(128, 64, 3),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(64, 64, 3),
            #nn.LeakyReLU(inplace=True),
        )
        self.final = nn.Conv2d(64, num_classes, 1)

    def forward(self, x):
        dec1 = self.dec1(x)
        dec2 = self.dec2(dec1)
        dec3 = self.dec3(dec2)
        dec4 = self.dec4(dec3)
        center = self.center(dec4)
        enc4 = self.enc4(torch.cat([
            center, F.upsample_bilinear(dec4, center.size()[2:])], 1))
        enc3 = self.enc3(torch.cat([
            enc4, F.upsample_bilinear(dec3, enc4.size()[2:])], 1))
        enc2 = self.enc2(torch.cat([
            enc3, F.upsample_bilinear(dec2, enc3.size()[2:])], 1))
        enc1 = self.enc1(torch.cat([
            enc2, F.upsample_bilinear(dec1, enc2.size()[2:])], 1))

        return F.upsample_bilinear(self.final(enc1), x.size()[2:])
</code></pre>
<p>and my optimizer</p>
<pre><code class="lang-auto">optimizer = optim.Adam(model.parameters(), lr=0.0001)
</code></pre>
<p>Also the output of loss</p>
<pre><code class="lang-auto">1 0.5320905093947355
2.0186073773985127
2 0.5344406446928762
1.8902393944431681
3 0.537645096729026
1.837425033883143
4 0.538817567193965
1.7639471042602735
5 0.5387050243204894
1.7177013124582658
</code></pre>
<p>Thank you if anyone know how can I improve the performance</p>

username: KFrank
cooked: <p>Hi Jason!</p>
<aside class="quote no-group quote-modified" data-username="Jason_Yuen" data-post="1" data-topic="136237" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jason_yuen/48/29871_2.png" class="avatar"> Jason_Yuen:</div>
<blockquote>
<pre><code class="lang-auto">        criterion = nn.CrossEntropyLoss()
        loss = criterion(output, target[:, 0])
...
        predicted = torch.argmax(output.data, 1)
</code></pre>
</blockquote>
</aside>
<p>This looks a little suspicious.  First you use <code>CrossEntropyLoss</code> to<br>
compare <code>output</code> with <code>target</code>, but with one dimension selected<br>
away.  But then in your accuracy calculation you perform the <code>.eq()</code><br>
test on <code>output</code> with one dimension reduced away by <code>.argmax()</code><br>
with <code>target</code> with all of its dimensions.</p>
<p>What are the shapes of <code>output</code> and <code>target</code>?  What are the meanings<br>
of their elements?  In particular, why are you selecting away one of<br>
<code>target</code>’s dimensions in your loss calculation?  (Is it a “channels” or<br>
“classes” dimension?)</p>
<p>Also, could you outline in a short sentence the problem you are<br>
working on?</p>
<aside class="quote no-group">
<blockquote>
<p>Also the output of loss</p>
<pre><code class="lang-auto">1 0.5320905093947355
2.0186073773985127
2 0.5344406446928762
1.8902393944431681
...
5 0.5387050243204894
1.7177013124582658
</code></pre>
</blockquote>
</aside>
<p>It’s hard to have much certainty without knowing the details of your<br>
use case, but it’s not at all clear that your training isn’t progressing<br>
reasonably.</p>
<p>Your loss is moving steadily down, and your accuracy is moving<br>
up, albeit slowly, with a slight downtick for the last epoch for which<br>
you gave us results.  This doesn’t seem so bad.</p>
<p>Also, for many realistic problems, five epochs is just getting started,<br>
and not really enough to draw conclusions about the training (but<br>
this depends on your use case and how large your training set is).<br>
What do the loss and accuracy curves look like if you train for, say,<br>
a hundred or several hundred epochs?</p>
<p>Best.</p>
<p>K. Frank</p>

username: Jason_Yuen
cooked: <p>Thank you for you reply.<br>
The shape of output, target and predicted are in the following</p>
<pre><code class="lang-auto">output torch.Size([1, 22, 256, 256])
#(Batch size, nClasses, H, W)
target torch.Size([1, 1, 256, 256])
#(Batch size, Channels, H, W)
predict torch.Size([1, 256, 256])
#(Batch size H, W)
</code></pre>
<aside class="quote no-group quote-modified" data-username="KFrank" data-post="2" data-topic="136237">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>In particular, why are you selecting away one of<br>
<code>target</code>’s dimensions in your loss calculation? (Is it a “channels” or<br>
“classes” dimension?)</p>
</blockquote>
</aside>
<p>I am dropping the channel dimension for loss calculation as I think that’s the required input type for CrossEntropyLoss and do you know the calculation method for accuracy is correct or not.</p>
<p>I am working on semantic segmentation problem and my training set is about 2000 images.</p>
<p>I found that the loss is keep dropping but a bit slow. I’m trying to  run for a larger epochs now</p>

username: KFrank
cooked: <p>Hi Jason!</p>
<aside class="quote no-group" data-username="Jason_Yuen" data-post="3" data-topic="136237" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jason_yuen/48/29871_2.png" class="avatar"> Jason_Yuen:</div>
<blockquote>
<p>The shape of output, target and predicted are in the following</p>
<pre><code class="lang-auto">output torch.Size([1, 22, 256, 256])
#(Batch size, nClasses, H, W)
target torch.Size([1, 1, 256, 256])
#(Batch size, Channels, H, W)
predict torch.Size([1, 256, 256])
#(Batch size H, W)
</code></pre>
</blockquote>
</aside>
<p>Yes, that should be okay then.  the “extra” dimension in <code>target</code> is<br>
a singleton (size-1) dimension, so, in some sense, it doesn’t really<br>
count.  (More precisely, when you “select it away” you don’t lose<br>
any information – the resulting tensor with one fewer dimension<br>
still has the same number of elements.)</p>
<aside class="quote no-group">
<blockquote>
<p>I am dropping the channel dimension for loss calculation as I think that’s the required input type for CrossEntropyLoss</p>
</blockquote>
</aside>
<p>Yes, this is correct.  You do need to get rid of <code>target</code>’s “extra” singleton<br>
dimension to keep <code>CrossEntropyLoss</code> from complaining.</p>
<aside class="quote no-group">
<blockquote>
<p>do you know the calculation method for accuracy is correct or not.</p>
</blockquote>
</aside>
<p>It looks correct.  (But don’t trust me – test and double-check it.)  Note<br>
that because <code>predicted</code> and <code>target</code> don’t have the same shape<br>
(<code>target</code> has that extra dimension), <em>broadcasting</em> will occur when you<br>
evaluate <code>predicted.eq (target.data)</code>.  But because it’s a singleton<br>
dimension, you’ll get the same result, so nothing will break.</p>
<aside class="quote no-group">
<blockquote>
<p>I am working on semantic segmentation problem and my training set is about 2000 images.</p>
</blockquote>
</aside>
<p>I haven’t used U-Net and I don’t really have any intuition about it,<br>
but I imagine that 2000 training images should be enough to train<br>
reasonably well.</p>
<p>I do notice that your accuracy starts off (that is, after just one epoch<br>
of training) at about 50%.  Given that you have 22 classes, this seems<br>
unexpectedly high.  Random guessing with 22 <em>uniformly distributed</em><br>
classes would give you an accuracy of only about 5%.</p>
<p>Is it possible that one of your classes is a “background” class and<br>
a large fraction of your pixels are background pixels?  For example,<br>
if half of your pixels were background pixels, you would get 50%<br>
accuracy by <em>always</em> predicting background.</p>
<p>If your data set is unbalanced because some classes are rare or<br>
especially common (such as in the background-pixel example),<br>
you might want to try <code>CrossEntropyLoss</code>’s class-weight <code>weight</code><br>
constructor argument to help counteract the class imbalance.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Jason_Yuen
cooked: <p>Thank you so much for your help. Some update about the training process, after 100 epoch the accuracy is about 80% so I think model is fine.</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="136237">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Is it possible that one of your classes is a “background” class and<br>
a large fraction of your pixels are background pixels?</p>
</blockquote>
</aside>
<p>Yes, the labels contains large amount of background. Thank you for your suggestion about the CrossEntropy with weight. I think it will be useful and I am trying it out.</p>