RuntimeError: weight tensor should be defined either for all or no classes

I am dealing with imbalance labels in my image segmentation. The model output -> 1, 5, 224, 224 is used for the prediction of two tasks task_a = output[1, [0, 1, 2], :,:] and task_b = output[1, [0, 1], :,:] . I am using the CrossEntropy loss and when I pass weight parameter `CrossEntropyLoss(weight = [0.5, 0.3, 0.1, 0.6, 0.5]), the runtime error is outputted

  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/pc-man/Documents/GAN/Code/main/losses/dice.py", line 84, in forward
    ce_loss = self.cross_entroy(predictions, labels)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 2266, in nll_loss
    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: weight tensor should be defined either for all or no classes

Could you check the shape of the model output again for this faulty batch as well as the shape and values of the target tensor?
Based on the error message I guess the shape of target values might have been changed (accidentally) during the training.

1 Like

It was because of the length of the class weight does not match with the number of classes i,e channels of the output and label.

Thanks for your contribution

Hi! may I ask how did you manage to add the background to your class weights?

I assumed my background will be 1’s for where the class label(s) are 0’s and otherwise. Then, I summed the 1’s in the background and other the class label(s) to get their respective frequency.

torch.manual_seed(19)

labelA = torch.randint(low=0, high=2, size=(2,3, 3), dtype=torch.float64)
labelB = torch.randint(low=0, high=2, size=(2,3, 3), dtype=torch.float64)

create_background = torch.zeros((labelB.shape), dtype=torch.long)

for label in [labelA, labelB]:
    create_background |= label.long()
create_background = create_background^1
print(labelA)
print(labelB)
print(create_background)

> tensor([[[1., 0., 1.],
         [0., 0., 1.],
         [0., 0., 1.]],

        [[0., 0., 0.],
         [0., 1., 1.],
         [0., 0., 0.]]], dtype=torch.float64)
tensor([[[0., 1., 1.],
         [1., 0., 1.],
         [1., 1., 1.]],

        [[1., 0., 1.],
         [1., 1., 0.],
         [1., 1., 0.]]], dtype=torch.float64)
tensor([[[0, 0, 0],
         [0, 1, 0],
         [0, 0, 0]],

        [[0, 1, 0],
         [0, 0, 0],
         [0, 0, 1]]])

class_frequency = [int(torch.sum(i)) for i in [create_background, labelA, labelB]]
print("Class weight: ", class_frequency)
> Class frequency:  [3, 6, 13]