Image segmentation using Unet to generate binary masks using pytorch and cuda

1

Hi, I am new to pytorch and image segmentation(UNET). I am trying to train the UNet to generate binary masks(Grayscale) on retinal images(RGB) + masks (RGB). I am converting both the masks to grayscale for training. Please help me with the following questions:

  1. The Mean IOU and Pixel Acc values are not getting computed. I am not entirely sure why? or if the function calls are active.
  2. How to implement pixel wise cross entropy loss? which loss function works efficient in this case?
  3. My training and validation loss values are very poor. Should I do anything different here?
  4. How should I save the model and use it for the testing loop.
    Thank you for your help.

Please find my code below.
Unet_Colab

2
  1. You can add debug print statements to these functions and check what is calculated and if there might be a bug in the code.
  2. You can use nn.CrossEntropyLoss with a model output shape of [batch_size, nb_classes, height, width] and a target in the shape [batch_size, height, width] containing class indices in the range [0, nb_classes-1].
  3. The accuracy and IOU looks a bit strange, so I would probably check them first.
  4. You can call model.eval() and wrap the testing loop into with torch.no_grad() to calculate the validation/testing results.
1 Like