Pytorch doing a cross entropy loss when the predictions already have probabilities

Question

So, normally categorical cross-entropy could be applied using a cross-entropy loss function in PyTorch or by combing a logsoftmax with the negative log likelyhood function such as follows:

m  = nn.LogSoftmax(dim=1)
loss = nn.NLLLoss()
pred   = torch.tensor([[-1,0,3,0,9,0,-7,0,5]], requires_grad=True, dtype=torch.float)
target = torch.tensor([4])
output = loss(m(pred), target)
print(output)

The thing is. What if the data at the output is already in a state with the probabilities where the variable pred already has the probabilities. Where the data is presented like the following:

pred = torch.tensor([[.25,0,0,0,.5,0,0,.25,0]], requires_grad=True, dtype=torch.float)

How could the cross-entropy then be completed in PyTorch?

Answer 1

You can implement categorical cross entropy pretty easily yourself. It is calculated as

$$ \text{cross-entropy} = -\frac{1}{n} \sum_{i=0}^{n} \sum_{j=0}^m \mathbf{y}_{ij} \log \hat{\mathbf{y}}_{ij} $$

where $n$ is the number of samples in your batch, $m$ is the number of classes, $\mathbf{y}_i$ is the one-hot target for example $i$, $\mathbf{\hat{y}}_i$ is the predicted probability distribution, and $\mathbf{y}_{ij}$ refers to the $j$-th element of this array.

In PyTorch:

def categorical_cross_entropy(y_pred, y_true):
    y_pred = torch.clamp(y_pred, 1e-9, 1 - 1e-9)
    return -(y_true * torch.log(y_pred)).sum(dim=1).mean()

You can then use categorical_cross_entropy just as you would NLLLoss in the training of a model. The reason that we have the torch.clamp line is to ensure that we have no zero elements, which will cause torch.log to produce nan or inf.

One difference you'll have to make in your code is that this version expects a one-hot target rather than an integer target. You can easily convert your current target list like so:

one_hot_targets = torch.eye(NUM_CLASSES)[targets]

where targets is a torch.tensor with integer values and NUM_CLASSES is the number of output classes that you have.