Change a semantic segmentation code from binary to multi-class，model predict error

PyTorch Live
1

I’m changing a semantic segmentation network from two-class to multi-class, and changed the last layer of the network to x=softmax(x,dim=1). Since it’s a six-class (no background class) problem, I changed num_classes to 6, and one-hot encoding is done on the targets to measure the loss. But when making predictions on the image, unclassified pixels appear on the image. Can someone help? Thanks in advance.

2

Could you share a small runnable snippet of code that highlights the problem? Make sure to include the loss function you used. Depending on the loss function used, it may or may not do the softmax for you already, and it may expect one-hot encoded target or an index target. Just going to guess that there’s some chance you have a mismatch between the loss function used and the rest of your model.

1 Like
3

Thanks for your hint, I changed the loss function to nn.CrossEntropy and removed the one-hot encoding of labels and the softmax function of the input in the network. After I check the colormap in plt.imsave(), the correct result showed. :laughing: :+1:

4

Delighted to hear it!