Loss does not decrease

vision
1

I have been trying modified U-Net on a custom dataset. I have attached the image and its corresponding mask. After some time the loss stops decreasing. It remains at around 1. The results produced are not satisfactory. The architecture I am using is :

class Unet(nn.Module):
    '''U-Net Architecture'''

    def __init__(self,inp,out):
        super(Unet,self).__init__()
        self.c1=self.contracting_block(inp,64)
        self.c2=self.contracting_block(64,128)
        self.c3=self.contracting_block(128,256)
        self.maxpool=nn.MaxPool2d(2)
        self.upsample=nn.Upsample(scale_factor=2,mode="bilinear",align_corners=True)
        self.c4=self.contracting_block(128+256,128)
        self.c5=self.contracting_block(64+128,64)
        self.c6=nn.Conv2d(64,out,1)

    def contracting_block(self,inp,out,k=3):
        block =nn.Sequential(
            nn.Conv2d(inp,out,k,padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(out),
            nn.Conv2d(out,out,k,padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(out)
        )
        return block


    def forward(self,x):
        #280,280,1
        conv1=self.c1(x) 
        #280,280,128
        x=self.maxpool(conv1)
        #140,140,128
        conv2=self.c2(x)
        #140,140,256
        x=self.maxpool(conv2)
        #70,70,256
        conv3=self.c3(x)
        #70,70,512
        x=self.upsample(conv3)
        #140,140,512
        x=torch.cat([conv2,x],axis=1)
        #140,140,768
        x=self.c4(x)
        #140,140,256
        x=self.upsample(x)
        #280,280,256
        x=torch.cat([conv1,x],axis=1)
        #280,280,384
        x=self.c5(x)
        #280,280,128
        x=self.c6(x)
        #280,280,3
        return x

I am using Adam as my optimizer and BCEWithLogitsLoss as my loss function. Any help would be highly appreciated.
Thank you.

2

I would try to scale down the problem a bit and try to overfit a small data sample first.
Once your model successfully overfits the small data sample, you could try to scale it up again.

Based on your loss function, it seems you are dealing with a multi-label segmentation use case, i.e. each pixel might belong to zero, one, or more classes.
If your use case is a multi-class segmentation, i.e. each pixel belongs to once class only, you would have to use nn.CrossEntropyLoss instead.

1 Like