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

moreshud: <p>I am dealing with imbalance labels in my image segmentation. The model output -&gt; <code>1, 5, 224, 224</code> is used for the prediction of two tasks <code>task_a = output[1, [0, 1, 2], :,:] and task_b = output[1, [0, 1], :,:] </code>. 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</p>
<pre><code class="lang-auto">  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
</code></pre>

ptrblck: <p>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?<br>
Based on the error message I guess the shape of target values might have been changed (accidentally) during the training.</p>

moreshud: <p>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.</p>
<p>Thanks for your contribution</p>

NCJ: <p>Hi! may I ask how did you manage to add the background to your class weights?</p>

moreshud: <p>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.</p>
<pre><code class="lang-auto">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)

&gt; 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)
&gt; Class frequency:  [3, 6, 13]

</code></pre>