Multiclass classification with softmax

I have 5 classes. I’m using a softmax function and getting 5 probabilities in each row that add up to 1 in total.
To calculate the loss do I just pick the column with the highest probability, assign it to the column number and then calculate loss. For eg: if column 1 has highest probability- that row is classified as Class 1 and let’s say the true output is Class 0. So the MSE loss is (1-0)^2 = 1. Or is there any better method?

Hi jpj!

Yes, you should use CrossEntropyLoss as your loss criterions (and not
use Softmax because CrossEntropyLoss expects logits rather than
probabilities as its predicted input).

Best.

K. Frank

1 Like