Which activation function to use in segmentation?

which activation function to be used at last layer of segmentation models like segnet, unet?
should I use F.sigmoid while defining the model’s last layer itself?

Hi Arjun!

The short answer is that you should just use the output of your last
linear layer, with no activation function, as input to your loss function
(for plain-vanilla use cases).

I assume that you are asking about training a network to segment
2d images, that is, the output of your network will be an “image”
whose pixel values tell you the segment that the original, input pixels
are in.

If you are doing binary segmentation, e.g., background / foreground,
the output of your last linear layer should have shape
[nBatch, height, width] (as should your target), and you should
use BCEWithLogitsLoss as your loss function.

For multi-class segmentation, e.g., background / normal tissue / benign
tumor / malignant tumor, the shape of your network’s output should be
[nBatch, nClass, height, width], your target should have shape
[nBatch, height, width], and you should use CrossEntropyLoss
as your loss function.

Again, in neither case should you feed the output of your last linear
layer through an activation function.


K. Frank

1 Like