Output not good, though loss is very less

vision
1

I have been trying to train a unet like model on a custom road segmentation dataset. You can look at an example here: sample_dataset
I am converting the mask to road lines using OpenCV edge detection.
The loss(both train and validation) during the training process became significantly low(1e-5). But the results during inference are not as expected(not even on the video from where I extracted the frames).
My dataLoader class looks like this

class load(Dataset):
    def __init__(self,**kwargs):
        self.width=kwargs["width"]
        self.height=kwargs["height"]
        self.samples=[]
        self.path1="/home/satinder/Desktop/deepWay/DeepWay.v2/dataSet/Segmentation2/img/"
        self.path2="/home/satinder/Desktop/deepWay/DeepWay.v2/dataSet/Segmentation2/mask/"
        img_folder=os.listdir(self.path1)
        
        for i in tqdm(img_folder):
            num=i.split(".")[0]
            self.samples.append((i,num+".png"))
        self.color=transforms.ColorJitter(brightness = 1)
        #self.translate=transforms.RandomAffine(translate=(0.1,0.1))
        self.angle=transforms.RandomAffine(degrees=(60))
        self.flip=transforms.RandomHorizontalFlip(p=0.5)
        self.transforms_img=transforms.Compose([transforms.ToTensor(),
                                                transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

        self.transforms_mask=transforms.Compose([transforms.Grayscale(num_output_channels=1),
                                                transforms.ToTensor(),
                                                transforms.Normalize((0.5,),(0.5,))])

    def __len__(self):
        return len(self.samples)

    def __getitem__(self,idx):
        i,j=self.samples[idx]
        img=cv2.imread(self.path1+i,1)
        img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        #img=cv2.blur(img,(3,3))
        mask=cv2.imread(self.path2+j,1)
        mask=cv2.cvtColor(mask,cv2.COLOR_BGR2GRAY)
        mask=cv2.Canny(mask,100,150)
        mask=cv2.dilate(mask,None,iterations=5)
        img=cv2.resize(img,(self.height,self.width))
        mask=cv2.resize(mask,(self.height,self.width))
        #print(mask.shape)
        seed=np.random.randint(2147483647)
        img=Image.fromarray(img)
        mask=Image.fromarray(mask)
        

        random.seed(seed)
        #img=self.color(img)
        random.seed(seed)
        #img=self.translate(img)
        random.seed(seed)
        #img=self.angle(img)
        random.seed(seed)
        #img=self.flip(img)
        random.seed(seed)
        img=self.transforms_img(img)
        
        random.seed(seed)
        #mask=self.translate(mask)
        random.seed(seed)
        #mask=self.angle(mask)
        random.seed(seed)
        #mask=self.flip(mask)
        random.seed(seed)
        mask=self.transforms_mask(mask)
        #print(img)
        return (img,mask)
    
    def plot(self,img):
        img=np.transpose(img.numpy(),(1,2,0))
        img=img*0.5+0.5
        img=cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
        cv2.imshow("ds",img)
        cv2.waitKey(0)


if(__name__=="__main__"):
    obj=load(width=256,height=256)
    res=obj.__getitem__(7)
    obj.plot(res[0])
    obj.plot(res[1])
    #cv2.imshow("img",res[0].cpu().detach().numpy())

I am using BCEWithLogitsLoss as my loss function and Adam as the optimizer.

2

How does the validation loss look during training?
Do you see any overfitting?

3

Sir validation loss also decreased with the training loss. I also considered the possibility of overfitting, I tried the model on the data I trained it on, then also the results were not as expected.

4

I’m not sure how your model works, but I assume it’s some kind of segmentation model?
If so, the model might overfit to the majority class, which is the background, as the white lines use very little spatial size compared to the background.
You could try to add weighting to your loss or something like focal loss.

1 Like
5

Sir, I am using a modified U-net like architecture: I am using this for a project of mine which runs on a embedded device, so I tried to make it as small as possible.

class Unet(nn.Module):
    '''U-Net Architecture'''
    def __init__(self,inp,out):
        super(Unet,self).__init__()
        self.c1=self.contracting_block(inp,8)
        self.c2=self.contracting_block(8,16)
        self.c3=self.contracting_block(16,32)
        self.maxpool=nn.MaxPool2d(2)
        self.upsample=nn.Upsample(scale_factor=2,mode="bilinear",align_corners=True)
        self.c4=self.contracting_block(16+32,16)
        self.c5=self.contracting_block(8+16,8)
        self.c6=nn.Conv2d(8,out,1)
        

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


    def forward(self,x):
        conv1=self.c1(x) 
        x=self.maxpool(conv1)
        conv2=self.c2(x)
        x=self.maxpool(conv2)
        conv3=self.c3(x)
        x=self.upsample(conv3)
        x=torch.cat([conv2,x],axis=1)
        x=self.c4(x)
        x=self.upsample(x)
        x=torch.cat([conv1,x],axis=1)
        x=self.c5(x)
        x=self.c6(x)
        #x=F.sigmoid(x)
        return x

Sir, the dataset I labelled had images like this
At that time also the the train and validation loss were very less, but the model wasn’t giving the desired results.
I thought maybe creating lanes would be beneficial, so I converted the masked images into images with lanes using opencv(edge detection and dilation operation) but turns out, the results have become much worse.
Sir, I am newbie can you point me to any resources for focal loss or how weight can be added to the loss ?
Thank you.