How to weight the loss?

for step, input in dataloader:
    outptut = model(input)  
    # top_prob, _ = torch.topk(F.softmax(output, dim=1)[0], 1)
    loss = criterion(output, label)

I want to weight the highest probability of softmax to the loss.
If, prob of softmax : 0.75 -> Effect of this sample * 0.75
If, prob of softmax : 0.12 -> Effect of this sample * 0.12

But the top_prob is the list (length == batch size)
like top_prob = [0.72, 0.75, 0.82, ... 0.94]
and loss is one scalar value.
(It means the loss is already computed with all elements of a batch.)

How to apply the probability of softmax to the loss ?

The CrossEntropy loss has a weight parameter for doing this, you can check it in the documentation.

Thanks for you answer.
But as far as I know, the weight in nn.CrossEntropyLoss() uses for the class-wise weight.
In my case, I need to weight sample-wise manner.

This is going to work:

import torch

x = torch.rand(16, 20)
y = torch.randint(2, (16,))
# Try torch.ones(16) here and it will be equivalent to
# regular CrossEntropyLoss
weights = torch.rand(16) 

net = torch.nn.Linear(20, 2)

def element_weighted_loss(y_hat, y, weights):
    m = torch.nn.LogSoftmax(dim=1)
    criterion = torch.nn.NLLLoss(reduction='none')
    loss = criterion(m(y_hat), y)
    loss = loss * weights
    return loss.sum() / weights.sum()

weighted_loss = element_weighted_loss(net(x), y, weights)

not_weighted_loss = torch.nn.CrossEntropyLoss()(net(x), y)

print(weighted_loss, not_weighted_loss)
1 Like

but weights.sum() is correct? not the size of batch?

Nope ! weights.sum() is there since the weights must sum up to 1.