Hello Ege!



 ekarais:

I am using one model to solve multiple classification tasks,

How can I use torch.nn.Softmax to achieve this?



First, for numerical-stability reasons, you shouldn’t use Softmax.
As I outline below, you should use CrossEntropyLoss, which has,
in effect, Softmax built into it.



How can I define the custom cross-entropy loss mentioned above?



You don’t need to write a custom cross-entropy loss.  Just use
pytorch’s built-in CrossEntropyLoss four times over, once for
each of your classification tasks.
Your model outputs a batch of prediction vectors of shape [nBatch, 22].
Your targets could be packaged in a number of ways.  The most
straightforward is probably to have four sets of targets, one for
each classification task.  Let’s call the four tasks A, B, C, and D.
Your targets should be batches of integer class labels.
So, for example, targetA should have shape [nBatch] and consist
of class labels that run from 0 to 4, because task A has five classes.
targetB should be the same.  targetC should also have shape
[nBatch], but consist of class labels that run from 0 to 7 because
task C has eight classes.
Then:
loss_fn = torch.nn.CrossEntropyLoss()  # only need to do this once

lossA = loss_fn (prediction[:,   0:5], targetA)
lossB = loss_fn (prediction[:,  5:10], targetB)
lossC = loss_fn (prediction[:, 10:18], targetC)
lossD = loss_fn (prediction[:, 18:22], targetD)

loss = lossA + lossB + lossC + lossD

That is, you use indexing to snip out of your vector of 22 predicted
class labels the set of predictions relevant to each task.  So in this
example, labels 10 through 17 (as indexed by 10:18) are the eight
predicted class labels relevant to task C.
Best.
K. Frank