I read the documentation of torch.nn.functional.softmax, but I am not very clear upon the usage of the
My output is of the following dimensions :
(batchsize, num_classes, length, breadth, height) - so it is a 5D output
I am doing a multi class semantic segmentation problem, so I want the pixel-wise softmax of corresponding pixels in
(length, breadth, height) along all classes. So do I use
dim = 1 while using
torch.nn.functional.softmax(output, dim = 1) , this way?
Thanks a lot!
torch.nn.functional.softmax(output, dim = 1) will yield an output, where the probabilities for each pixel sum to 1 in the class dimension (dim1).
This probability tensor can be used as a sanity check or for visualization purposes.
However, note that e.g.
nn.CrossEntropyLoss expects raw logits as the model’s output, since internally
F.log_softmax(output, 1) and
nn.NLLLoss will be used, so that you shouldn’t use the softmax method manually.
Alright, thanks a lot!