How to visualize segmentation result with CrossEntropyLoss

Question

I know this question already has a clear answer about how to visualize segmentation result. My question is if I'm using the CrossEntropyLoss as the loss function, after training. The predicted classification (segmentation) result doesn't represent the class 0,1,2. It represents the unnormalized Log probability as told in cs231n standford class. So if I want to visualize the segmentaion result, should I use the

- nn.LogSoftmax()

to get the classes and then plot them out? or should I do a maximum function over the first dimension of my output.

For example, if the predicted output is size 3x1024x1024, how could I visualize the segmenation result?

Thank you in advance.

Answer 1

In either case (with or without using nn.LogSoftMax()) you should do some argmax operation to get the class from probabilities.

The reason for that is that order of probabilities will always stay the same after the nn.LogSoftmax() since that operation squeezes the numbers in the [0, 1] interval, but keeps the ordering (biggest one before this operation will also remain the biggest in the [0, 1] interval, and so on ...).

Hence, for the purpose of determinating the class, it doesn't matter whether you do nn.LogSoftmax() or not.

nn.LogSoftmax() is needed only during the training process in order to calculate the loss and backpropagate that information to update the parameters of the network.