Multi-class cross entropy loss and softmax in pytorch

[KFrank]

Hi Brando!

@Brando_Miranda

What loss function are we supposed to use when we use the F.softmax layer?

If you want to use a cross-entropy-like loss function, you shouldn’t
use a softmax layer because of the well-known problem of increased
risk of overflow.

I gave a few words of explanation about this problem in a reply in
another thread:

Softmax regression from scratch: NaN loss

Hi Lorenzo! I didn’t look at your code, but if you wrote your softmax and cross-entropy functions as two separate functions you are probably tripping over the following problem. Softmax contains exp() and cross-entropy contains log(), so this can happen: large number --> exp() --> overflow NaN --> log() --> still NaN even though, mathematically (i.e., without overflow), log (exp (large number)) = large number (no NaN). Pytorch’s CrossEntropyLoss (for example) uses standard techniques …

You should either use nn.CrossEntropyLoss (which takes
pre-softmax logits, rather than post-softmax probabilities)
without a softmax-like layer, or use a nn.LogSoftmax layer,
and feed the results into nn.NLLLoss. (Both of these combine
an implicit softmax with the subsequent log in a way that avoids
the enhanced overflow problem.)

If you are stuck for some reason with your softmax layer, you
should run the probabilities output by softmax through log(),
and then feed the log-probabilities to nn.NLLLoss (but expect
increased risk of overflow).

(I am not aware of any single pytorch cross-entropy loss function
that takes post-softmax probabilities directly.)

Good luck!

K. Frank