BCEWithLogitsLoss Error on Batch Images

Rasoul: <p>I don’t understand why BCEWithLogitsLoss doesn’t work on batch images. Here is a minimum working example to reproduce the error,</p>
<pre><code class="lang-auto">import torch
target = torch.ones([10, 64, 20, 20], dtype=torch.float32)  # 64 classes, batch size = 10
output = torch.full([10, 64, 20, 20], 1.5)  # A prediction (logit)
pos_weight = torch.ones([64])  # All weights are equal to 1
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
criterion(output, target)
</code></pre>
<p>results in the error message,<br>
<code>RuntimeError: The size of tensor a (64) must match the size of tensor b (20) at non-singleton dimension 3</code></p>

KFrank: <p>Hi Rasoul!</p>
<aside class="quote no-group" data-username="Rasoul" data-post="1" data-topic="154492" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rasoul/48/37728_2.png" class="avatar"> Rasoul:</div>
<blockquote>
<p>I don’t understand why BCEWithLogitsLoss doesn’t work on batch images.</p>
</blockquote>
</aside>
<p><code>BCEWithLogitsLoss</code> does work on batch* images.</p>
<aside class="quote no-group">
<blockquote>
<p>Here is a minimum working example to reproduce the error,</p>
<pre><code class="lang-auto">target = torch.ones([10, 64, 20, 20], dtype=torch.float32)  # 64 classes, batch size = 10
output = torch.full([10, 64, 20, 20], 1.5)  # A prediction (logit)
pos_weight = torch.ones([64])  # All weights are equal to 1
</code></pre>
<p><code>RuntimeError: The size of tensor a (64) must match the size of tensor b (20) at non-singleton dimension 3</code></p>
</blockquote>
</aside>
<p>The short story is that <code>pos_weight</code> has to <a href="https://pytorch.org/docs/1.11/notes/broadcasting.html" rel="noopener nofollow ugc">broadcast</a> with respect to <code>output</code><br>
and <code>target</code>, and in your case it doesn’t.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.11.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; target = torch.ones([10, 64, 20, 20], dtype=torch.float32)  # 64 classes, batch size = 10
&gt;&gt;&gt; output = torch.full([10, 64, 20, 20], 1.5)  # A prediction (logit)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss() (output, target)
tensor(0.2014)
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = torch.ones ([64, 1, 1])) (output, target)
tensor(0.2014)
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = torch.ones ([64, 20, 1])) (output, target)
tensor(0.2014)
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = torch.ones ([64, 20, 20])) (output, target)
tensor(0.2014)
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = torch.tensor ([1.])) (output, target)
tensor(0.2014)
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = torch.ones ([64])) (output, target)
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "&lt;path_to_pytorch_install&gt;\torch\nn\modules\module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "&lt;path_to_pytorch_install&gt;\torch\nn\modules\loss.py", line 713, in forward
    return F.binary_cross_entropy_with_logits(input, target,
  File "&lt;path_to_pytorch_install&gt;\torch\nn\functional.py", line 3132, in binary_cross_entropy_with_logits
    return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
RuntimeError: The size of tensor a (64) must match the size of tensor b (20) at non-singleton dimension 3
</code></pre>
<p>*)  The <a href="https://pytorch.org/docs/1.11/generated/torch.nn.BCEWithLogitsLoss.html" rel="noopener nofollow ugc">BCEWithLogitsLoss documentation</a> is quite elliptical on this issue,<br>
but, as near as I can tell (from experiments rather than the documentation),<br>
<code>BCEWithLogitsLoss</code> doesn’t actually have any notion of “batch size”<br>
or “number of classes” or assign any particular meaning to any of the<br>
dimensions.  Rather, <code>output</code>, <code>target</code>, <code>weight</code>, and <code>pos_weight</code> all<br>
have to broadcast with respect to one another.  It then just performs an<br>
element-wise loss computation which gets reduced.</p>
<p>Best.</p>
<p>K. Frank</p>

Rasoul: <p>Hi Frank,</p>
<p>Thanks for the explanation!</p>
<p>The documentatuon, as you mentioned is not clear in this regard.</p>
<p>Best,</p>
<p>Rasoul</p>