Hi all,
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 (https://www.kaggle.com/cankls3/carvana-pytorch 6).
I have a training  dataset of 600  RGB images (224x224) labeled that I’m passing to the model:
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

Followed by:
model = UNet(n_class=1)
model = model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-5)
criterion = nn.BCELoss()
model.train()

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.
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 ).
Any suggestion in terms of structure of the model or parameters to tune? Maybe the training set is not enough big.
This is my training for loop:
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))

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
Thanks in advance for any help