Pytorch Equivalent To Tf.nn.softmax_cross_entropy_with_logits And Tf.nn.sigmoid_cross_entropy_with_logits

I found the post here. Here, we try to find an equivalence of tf.nn.softmax_cross_entropy_with_logits in PyTorch. The answer is still confusing to me. Here is the Tensorflow 2 code

Solution 1:

- tf.nn.softmax_cross_entropy_with_logits

Edit: This is actually not equivalent to F.cross_entropy. The latter can only handle the single-class classification setting. Not the more general case of multi-class classification, whereby the label can be comprised of multiple classes. Indeed, F.cross_entropy takes a unique class id as target (per instance), not a probability distribution over classes as tf.nn.softmax_cross_entropy_with_logits can expect to receive.

>>> logits = torch.tensor([[4.0, 2.0, 1.0], [0.0, 5.0, 1.0]])
>>> labels = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.8, 0.2]])

In order to get the desired result apply a log-softmax to your logits then take the negative log-likelihood:

>>>-torch.sum(F.log_softmax(logits, dim=1) * labels, dim=1)
tensor([0.1698, 0.8247])

---

- tf.nn.sigmoid_cross_entropy_with_logits

For this one you can apply F.binary_cross_entropy_with_logits.

>>> F.binary_cross_entropy_with_logits(logits, labels, reduction='none')
tensor([[0.0181, 2.1269, 1.3133],
        [0.6931, 1.0067, 1.1133]])

It is equivalent to applying a sigmoid then the negative log-likelihood, considering each class as a binary classification task:

>>> labels*-torch.log(torch.sigmoid(logits)) + (1-labels)*-torch.log(1-torch.sigmoid(logits))
tensor([[0.0181, 2.1269, 1.3133],
        [0.6931, 1.0067, 1.1133]])

---

having imported torch.nn.functional as F.