Inconsistency CrossEntropyLoss vs BCELoss regarding logits/probability space

[ego-thales]

🚀 The feature, motivation and pitch

The loss CrossEntropyLoss expects logits inputs. On the other hand, there exist BCELoss and BCEWithLogitsLoss. But there is no counterpart to CrossEntropyLoss that expects inputs in the probability space.

Alternatives

I think it would make more sense if CrossEntropyLoss expected probability inputs and if a new CrossEntropyWithLogitsLoss was introduced. But this would be a dramatic change in the API.

As such, I would propose to introduce an option with_logits (that defaults to True to keep compatibility) to CrossEntropyLoss.

What do you think?

Additional context

Finally, I ask the following questions:

• Is it possible to easily train a neural net that outputs softmax probabilities and not logits? It seems like since there is no CrossEntropyLoss counterpart for probability space, the user is forced to add a extra step and use NLLLoss or write a custom loss.
• Is the fact that CrossEntropyLoss doesn't work with probability space due to under/overflow handling in general?

Thanks!

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

[rajveer43]

Can I Work on it?

[jbschlosser]

Note that nn.NLLLoss takes in inputs as log-probabilities, and operating on these is generally more numerically stable.

Is the fact that CrossEntropyLoss doesn't work with probability space due to under/overflow handling in general?

So I'd say yes to this :)