Hi Droid!



 droid:

I have a vanilla implementation of UNet, which I want to use for multiclass segmentation (where each pixel can belong to many classes).


Let me distinguish between a (single-label) multi-class problem and
a multi-label, multi-class problem.  In both cases you have multiple
classes (one of which might be a “background” or catch-all “other”
class).
In the single-label case, each pixel (or more generally, each item
to be classified) belongs to exactly one class.  CrossEntropyLoss
is the natural loss criterion in this case.
In the multi-label case, each pixel belongs to any number of the
classes, including no class or all classes.  BCEWithLogitsLoss
is the natural loss criterion to use for a multi-label, multi-class
problem, and it works out of the box for this use case without any
additional fussing around.
You say that “each pixel can belong to many classes,” so I understand
you to be working with a multi-label, multi-class problem.
The multi-label problem should be understood as number-of-classes
binary classification problems all being run in parallel in the same
network.  Is the pixel in class-A, yes or no? (a binary problem);
independently, is the pixel in class-B, yes or no? (a second binary
problem); and so on.
(Please note that BCEWithLogitsLoss is to be preferred to BCELoss
for reasons of numerical stability.)
Best.
K. Frank