Hi Droid!



 droid:

I do not have a background class – just 3 output channels for R, G, B.


First, a clarifying question about your multi-class segmentation
problem:
In the “Mask” (and “Predicted”) image you post I see four colors:
black, green, blue, and magenta.  You might have other images
with more colors, but it looks to me like you might be performing
single-label, four-class (or perhaps a few more than four-class)
segmentation.
I understand that magenta is made of red and blue, so you could
interpret it as being in the both the “red” class and the “blue” class,
but is there any substantive sense in which the pixels making up
the QR code in your “Image” image are logically both in a “red”
and “blue” class, or did you just happen to label the QR pixels with
magenta in your mask?
Just because you choose to label pixels in your mask with saturated
RGB colors doesn’t mean that your problem is naturally a multi-label,
three-class (RGB) problem.  Is it possible that single-label, multi-class
segmentation is a better logical fit for your use case?


I am completely missing one of my output channels (the green channel in the below example). It seems like because this class is small in magnitude, compared to the other classes, the network is not learning this class at all.


Whether you treat this as single-label or multi-label, the fact that one
of your classes (green) is significantly less common than the others
can lead to an otherwise well-trained network not predicting the “rare”
class at all.
Class weights, where you weight the green class more heavily, make
sense in such a case, and both CrossEntropyLoss (for the single-label
case) and BCEWithLogitsLoss (for the multi-label case) support class
weights.  A common, sensible approach is to weight each class in
inverse proportion to how often it occurs in your training set.


I create the predicted output image using torch.round(F.sigmoid(model(data))) to force every pixel to 0 or 1 for each class.


As a side note, you don’t need to use sigmoid() / round() to convert
your floating-point network output to a binary prediction.  Thresholding
against zero, data > 0.0, is equivalent, and I prefer it stylistically.
Best.
K. Frank