Issue with nn.CrossEntropyLoss() Segmentation

danishnazir: <p>Hi I have a prediction of shape [1,1,136,100]  where Class is one and Batch size is 1.<br>
My Label  is of size [1,136,100]<br>
Label = [[[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,<br>
0,   0,255,…]]]<br>
I am recieving cuda assert error.<br>
Further Trace is given as<br>
self.criterion = nn.CrossEntropyLoss(ignore_index=255)<br>
loss_semantic_seg = self.criterion(mask_pred, labels)<br>
File “/home/danish/miniconda3/envs/open-mmlab/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 532, in <strong>call</strong><br>
result = self.forward(*input, **kwargs)<br>
File “/home/danish/miniconda3/envs/open-mmlab/lib/python3.7/site-packages/torch/nn/modules/loss.py”, line 916, in forward<br>
ignore_index=self.ignore_index, reduction=self.reduction)<br>
File “/home/danish/miniconda3/envs/open-mmlab/lib/python3.7/site-packages/torch/nn/functional.py”, line 2021, in cross_entropy<br>
return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)<br>
File “/home/danish/miniconda3/envs/open-mmlab/lib/python3.7/site-packages/torch/nn/functional.py”, line 1840, in nll_loss<br>
ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)</p>

KFrank: <p>Hi Danish!</p>
<aside class="quote no-group" data-username="danishnazir" data-post="1" data-topic="116475" 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/danishnazir/48/33515_2.png" class="avatar"> danishnazir:</div>
<blockquote>
<p>I have a prediction of shape [1,1,136,100]  where Class is one and Batch size is 1.<br>
My Label  is of size [1,136,100]<br>
Label = [[[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,<br>
0,   0,255,…]]]</p>
</blockquote>
</aside>
<p>Based on this I speculate that you are performing binary segmentation<br>
of an image.  That is, for each pixel in your sample image, you want to<br>
predict whether it is “background” or “foreground.”</p>
<p>If so, you should use be using <code>BCEWithLogitsLoss</code>.</p>
<p>Let’s say that your sample images have shape <code>[height, width]</code>.  Then<br>
an input batch will have shape <code>[nBatch, height, width]</code>, and your<br>
prediction will have the <em>same</em> shape, with no class dimension.  So,<br>
using your sizes, both your input batch and prediction will have shape<br>
<code>[1, 136, 100]</code>.</p>
<p>Your <code>target</code> (your “Label”) should also have this same shape, and<br>
contain values in the range <code>[0.0, 1.0]</code>, inclusive.  For a purely binary<br>
<code>target</code> (e.g., black-and-white), these values can be exactly <code>0.0</code> and<br>
<code>1.0</code>.  Note that for a black-and-white segmentation mask you do not<br>
want values of <code>0</code> and <code>255</code>, but rather <code>0.0</code> and <code>1.0</code>.</p>
<p>Lastly, your prediction should be the raw-score logits output by your<br>
final <code>Linear</code> layer, and not be passed through a <code>sigmoid()</code> that would<br>
convert them to probabilities.</p>
<p>Best.</p>
<p>K. Frank</p>

danishnazir: <p>Thanks Frank, I think I have found the issue and it is because of the label range which is from [0,255]. Can you tell me how I can convert it to [0,1]? based on some threshold e.g 127.</p>

KFrank: <p>Hi Danish!</p>
<aside class="quote no-group" data-username="danishnazir" data-post="3" data-topic="116475" 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/danishnazir/48/33515_2.png" class="avatar"> danishnazir:</div>
<blockquote>
<p>Can you tell me how I can convert it to [0,1]? based on some threshold e.g 127.</p>
</blockquote>
</aside>
<p>Thresholding will certainly work.  You could also use <code>torch.where()</code>.</p>
<p>If you <em>know</em> that your values will always be exactly <code>0</code> or <code>255</code>, you<br>
could divide by <code>255</code> (and convert to float):</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; targ = torch.tensor ([0, 255, 0, 0, 255])
&gt;&gt;&gt; targ.dtype
torch.int64
&gt;&gt;&gt; targ = (targ / 255).float()
&gt;&gt;&gt; targ
tensor([0., 1., 0., 0., 1.])
&gt;&gt;&gt; targ.dtype
torch.float32
</code></pre>
<p>Best.</p>
<p>K. Frank</p>