U-net segmentation loss does not decrease

username: Diauro
cooked: <p>Hi all,</p>
<p>I’m quite new to pytorch so I know I could be doing some very basic mistakes. Currently I’m trying to do image segmentation using U-net using a structure of notebook very similar to this notebook of the carvana-segmentation competition in kagglem (<a href="https://www.kaggle.com/cankls3/carvana-pytorch" rel="nofollow noopener">https://www.kaggle.com/cankls3/carvana-pytorch</a>).</p>
<p>I have a training  dataset of 600  RGB images (224x224) labeled that I’m passing to the model:</p>
<pre><code class="lang-auto">train_data_loader = DL.DataLoader(train_dataset, batch_size=1, shuffle=False)


def double_conv(in_channels, out_channels):

    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, 3, padding=1),
        nn.ReLU(inplace=True)
    )   

class UNet(nn.Module):

    def __init__(self, n_class):
        super().__init__()
                
        self.dconv_down1 = double_conv(3, 64)
        self.dconv_down2 = double_conv(64, 128)
        self.dconv_down3 = double_conv(128, 256)
        self.dconv_down4 = double_conv(256, 512)  

        self.maxpool = nn.MaxPool2d(2)
        
        self.dconv_up3 = double_conv(512, 256)
        self.dconv_up2 = double_conv(256, 128)
        self.dconv_up1 = double_conv(128, 64)
        
        self.TConv3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.TConv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.TConv1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        
        self.conv_last = nn.Conv2d(64, n_class, 1)
        
    def forward(self, x):
        conv1 = self.dconv_down1(x)
        x = self.maxpool(conv1)
        
        conv2 = self.dconv_down2(x)
        x = self.maxpool(conv2)
        
        conv3 = self.dconv_down3(x)
        x = self.maxpool(conv3)
        
        x = self.dconv_down4(x)
        
        x = self.TConv3(x)
        x = torch.cat([x, conv3], dim=1)

        x = self.dconv_up3(x)
        x = self.TConv2(x)
        x = torch.cat([x, conv2], dim=1)

        x = self.dconv_up2(x)
        x = self.TConv1(x)
        x = torch.cat([x, conv1], dim=1)

        x = self.dconv_up1(x)
        out = self.conv_last(x)
        out = F.sigmoid(out)
        
        return out
</code></pre>
<p>Followed by:</p>
<pre><code class="lang-auto">model = UNet(n_class=1)
model = model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-5)
criterion = nn.BCELoss()
model.train()
</code></pre>
<p>When I start to train in the first 20 epochs the loss goes down quite fast but then starts to saturate around 0.63 and does not really moves from these values.</p>
<p>I’ve tried changing the learning rate and tried to use SGD optimizer but I did not have any significat result (actually with SGD the loss increases ).</p>
<p>Any suggestion in terms of structure of the model or parameters to tune? Maybe the training set is not enough big.</p>
<p>This is my training for loop:</p>
<pre><code class="lang-auto">model.train()

num_epochs = 50
running_loss = 0.0

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (X, y) in enumerate(train_data_loader):
        X = X.to(device)
        y = y.to(device)
        X = Variable(X)
        y = Variable(y)
        
        optimizer.zero_grad()
        output = model(X)
        loss = criterion(output, y)
        
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        
        
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            } 
        torch.save(checkpoint, 'checkpoint.pth')
    
    
        # accuracy
        #_, predicted = torch.max(output, y)
        #total_train += y.nelement()
        #correct_train += predicted.eq(y.data).sum().item()
        #train_accuracy = 100 * correct_train / total_train
        #avg_accuracy = train_accuracy / len(train_loader)  
        #,  "Training Accuracy: %d %%" % (train_accuracy)
    
    print("loss for epoch " + str(epoch) + ":  " + str(running_loss))
</code></pre>
<p>P.s I’m trying to implement the train accuracy but for the moment is not working so I’ve commented those lines, in case you have ideas please don’t esistate</p>
<p>Thanks in advance for any help</p>

username: ptrblck
cooked: <p>The code looks generally alright.<br>
You could try to remove the sigmoid activation at the end of your model and instead use <code>nn.BCEWithLogitsLoss</code>. This would give you more numerical stability and might avoid vanishing gradients if your output saturates.</p>
<p>Besides that, you should remove the <code>Variable</code> calls, as they are deprecated since <code>0.4.0</code>.<br>
Also, the usage of <code>.data</code> is not recommended, as this might yield silent errors, since Autogradd cannot track these operations.</p>

username: Diauro
cooked: <p>Thanks your reply!</p>
<p>So basically it would comment <code>out = F.sigmoid(out)</code> and placing the loss with<br>
<code>out = nn.BCEWithLogitsLoss(out)</code> ? .</p>
<p>Also I was considering that the very low decrease of the loss is due to my images that are RGB but not normalized (like to resnet values or with their std and avg).</p>
<p>Currently I’m loading my data by using the following code:</p>
<p>class MyDataset_test(Dataset):<br>
def <strong>init</strong>(self, image_paths ,train=True):<br>
self.image_paths = image_paths<br>
<span class="hashtag">#self</span>.mask_paths = mask_paths</p>
<pre><code class="lang-auto">    def transforms(self, image, mask):
        #img = img.resize((wsize, baseheight), PIL.Image.ANTIALIAS)
        #image = transforms.Resize(size=(64, 64))(image)
        #mask = transforms.Resize(size=(64, 64))(mask)
        image = image.resize((64, 64), PIL.Image.NEAREST)
        mask = mask.resize((64, 64), PIL.Image.NEAREST)
        image = TF.to_tensor(image)
        mask = TF.to_tensor(mask)
        
        return  mask
    
    def __getitem__(self, index):
        #image = Image.open(self.image_paths[index])
        mask = Image.open(self.mask_paths[index])
        
        x, y = self.transforms(mask)
        return x, y
    
    def __len__(self):
        return len(self.image_paths)
</code></pre>
<p>maybe this could be a possible issue.</p>
<p>Thanks in advance!</p>

username: ptrblck
cooked: <p>Use <code>nn.BCEWithLogitsLoss</code> as the criterion instead of <code>nn.BCELoss</code> and pass the raw logits to it by removing the sigmoid.</p>
<p>Normalization might help your training so you should definitely try it out.<br>
Also you could try to overfit a small data sample and make sure your current training routine works fine.<br>
Once this works you could try to scale it up again.</p>

username: Diauro
cooked: <p>Currently I’m using as you suggested <strong>nn.BCEWithLogitsLoss</strong> and tried the model on a small data set of 15 images (should be small enough). But now the loss even if it goes down very fast it  saturates and  around 0.086 and 0.084.<br>
So I guess this indicates that there is a more serious problem on the model or on the data?</p>
<p>I’ve tried out of curiosity this example using ResNetUnet from this example <a href="https://github.com/usuyama/pytorch-unet" rel="nofollow noopener">ResNetUnet</a> and again the loss saturates for a small training data around 0.083.</p>
<p>I’ve also normalized my data with</p>
<pre><code class="lang-auto">transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
</code></pre>
<p>Any suggestion is more than welcome. Thanks again !</p>

username: Diauro
cooked: <p>I was looking to this post (<a href="https://discuss.pytorch.org/t/unet-implementation/426/23" class="inline-onebox">UNet implementation</a> a bit old) where apparently in pytorch there were some issues to implement Unet.  I could not find a real solution for the moment, Saed in one reply wrote only</p>
<p>"For the last set of convolutions, that is 128-&gt; 64 -&gt; 64 -&gt; 1, the activation function should not be used!<br>
The activation function causes the values to vanish!</p>
<p>I just removed the nn.ReLU() modules on top of these convolution layers and now everything works like a charm! "</p>