Hi Chris!



 Chris_XU:

But this involves looping which I assume would be slower than a direct vectorized computation.


You can get a tensor of weights corresponding to the class labels
in t by indexing into w, and you can get a tensor of (negative)
log-probabilities by calling .take_along_dim() with t as the
argument:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> _ = torch.manual_seed (2021)
>>> K = 10
>>> X = torch.randn(32, K, 20, 30)
>>> t = torch.randint(K, (32, 20, 30)).long()
>>> w = torch.randn(K).abs()
>>> xeA = torch.nn.CrossEntropyLoss (weight = w, reduction = 'none') (X, t)
>>> xeB = w[t] * -X.log_softmax (1).take_along_dim (t.unsqueeze (1), 1).squeeze (1)
>>> torch.allclose (xeA, xeB)
True
>>> torch.equal (xeA, xeB)
True

As an aside:


    loss = -w[label] * torch.log(XX[ind][label].exp() / XX[ind].exp().sum())



you are computing log-softmax() “by hand” in the straightforward way,
but it is numerically more stable to use the log-sum-exp trick, either
by implementing it explicitly or by using pytorch’s log_softmax()
function (as I have done above).
Best.
K. Frank