BCEWithLogits with weight/pos_weight

islomjon: <p>Hello, I have read several topics about setting the weight to a loss, but I have some interesting to me question.</p>
<p>So I have a binary segmentation problem, with classes 0 – background, and 1 – buildings. And they are unbalanced. I decided to set a weight for BCEWithLogits loss with <code>torch.tensor([0.3, 0.7])</code> for class 0 and 1, respectively. But when I try to calculate the loss, script throws a Runtime error:</p>
<pre><code class="lang-auto">import torch
output = torch.randn(1, 2, 256, 256, requires_grad=True)
target = torch.randn(1, 2, 256, 256, requires_grad=False)
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([0.3, 0.7])) 
# same when I use 'weight' insted pos_weight
loss = criterion(output, target)
</code></pre>
<p><strong>Error</strong></p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-9-0ea142ac673d&gt; in &lt;module&gt;
      3 target = torch.randn(1, 2, 256, 256, requires_grad=False)
      4 criterion = torch.nn.BCEWithLogitsLoss(weight=torch.tensor([0.3, 0.7]))
----&gt; 5 loss = criterion(output, target)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--&gt; 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    615                                                   self.weight,
    616                                                   pos_weight=self.pos_weight,
--&gt; 617                                                   reduction=self.reduction)
    618 
    619 

/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   2433         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   2434 
-&gt; 2435     return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
   2436 
   2437 

RuntimeError: The size of tensor a (256) must match the size of tensor b (2) at non-singleton dimension 3
</code></pre>
<p>But when I try only one tensor with shape 1, it turns out that it works</p>
<pre><code class="lang-auto">import torch
output = torch.randn(1, 2, 256, 256, requires_grad=True)
target = torch.randn(1, 2, 256, 256, requires_grad=False)
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([0.3]))
loss = criterion(output, target)

print(loss)

&gt;&gt;&gt; tensor(0.2424, grad_fn=&lt;BinaryCrossEntropyWithLogitsBackward&gt;)
</code></pre>
<p>My question is why the first method didn’t work, and the second worked. And if the second method is the correct one, then which class will get correct weight to get better loss?</p>
<p>Thank you!</p>

ptrblck: <aside class="quote no-group" data-username="islomjon" data-post="1" data-topic="160298">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/i/d78d45/48.png" class="avatar"> islomjon:</div>
<blockquote>
<p>So I have a binary segmentation problem, with classes 0 – background, and 1 – buildings.</p>
</blockquote>
</aside>
<p>A binary segmentation use case with <code>nn.BCEWithLogitsLoss</code> would expect the model output and the target to have the shape <code>[batch_size, 1, height, weight]</code> where a low logit would indicate class0 and a high logit would indicate class1.<br>
For this use case the <code>pos_weight</code> should also be a single value as described in <a href="https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html#torch.nn.BCEWithLogitsLoss">the docs</a>.</p>
<p>The error message is still surprising and I would guess PyTorch tries to broadcast the <code>pos_weight</code> in your code. For a multi-class segmentation, where each pixel can have 0, 1, or multiple active classes you might need to <code>unsqueeze</code> the <code>pos_weight</code> in <code>dim2</code> and <code>dim3</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">output = torch.randn(1, 2, 256, 256, requires_grad=True)
target = torch.randn(1, 2, 256, 256, requires_grad=False)
pos_weight = torch.tensor([0.3, 0.7])
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight[None, :, None, None])
loss = criterion(output, target)
</code></pre>