Hi Mohsen!



 moh_sen:

I want to compute sum of cross entropy over all classes for each prediction
def sum_of_CE_lost(input):
    L = 0
    for c in range(1000):
        L = L + torch.nn.CrossEntropyLoss(input, c)
    return L

However, it is very slow. What is a better way? How can we parallelized it for GPU (CUDA)?


Yes, this is straightforward, based on the definition of cross_entropy.
Let input be a tensor of shape (nBatch, nClass).  Then (mentally
ignoring the minor errors in your code such as using the class version
of CrossEntropyLoss as if it were a function and c not being a tensor)
the following will match the result of the code you posted:
-torch.nn.functional.log_softmax (input).sum() / input.shape[0]

(We divide by input.shape[0] because cross_entropy() takes, by
default the mean across the batch dimension.)
Because this expression uses pytorch tensor functions, you will
automatically get the benefit of pytorch’s gpu support (if you move
your tensors to the gpu) (as well as autograd, if you care).
(As an aside, though, it’s not really clear to me that your calculation
makes a lot of sense.  But you can compute it with pytorch tensors
on the gpu without explicit loops.)
Good luck!
K. Frank