Pytorch softmax classificatoin

Hi everyone,

I am a little bit confused now, I have a small toy problem in binary classification and I am use CrossEntropyLoss() as an criterion without attaching LogSoftmax() at the end of the architecture, since it was written that the loss is calculating it implicitly. The question now, during testing, should I add LogSoftmax() or Softmax() ?

It depends which information and output you want to create during testing.
If you are only interested in the predicted classes you can directly use preds = torch.argmax(output, dim=1) on the raw logits, since the order won’t change if (log)softmax is applied.
On the other hand, if you want to see the probabilities for some reason you could apply a softmax to the output. Note however, that the loss calculation via nn.CrossEntropyLoss would still expect to receive raw logits.

1 Like