CrossEntropy loss function error - target 3 is out of bounds

benkolber: <p>Hey There!</p>
<p>I’m running a semantic segmentation model (unet) on input images of shape (3, 256, 256) and masks of shape (256,256) with pixel values of 0,1,2, and 3 (3 classes total, with pixel 0 being the background).</p>
<p>When retrieving the output from the model, I have:<br>
output from model = (b_s, 3, 256, 256)<br>
predicted mask = (b_s, 256, 256)</p>
<p>I use a cross-entropy loss function, but receive the following error:</p>
<pre><code class="lang-auto">   1838         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   1839     elif dim == 4:
-&gt; 1840         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   1841     else:
   1842         # dim == 3 or dim &gt; 4

IndexError: Target 3 is out of bounds.
</code></pre>
<p>This is the model I am using: <a href="https://github.com/milesial/Pytorch-UNet" rel="noopener nofollow ugc">source</a></p>
<pre><code class="lang-auto">&gt; class DoubleConv(nn.Module):
&gt;     """(convolution =&gt; [BN] =&gt; ReLU) * 2"""
&gt; 
&gt;     def __init__(self, in_channels, out_channels, mid_channels=None):
&gt;         super().__init__()
&gt;         if not mid_channels:
&gt;             mid_channels = out_channels
&gt;         self.double_conv = nn.Sequential(
&gt;             nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
&gt;             nn.BatchNorm2d(mid_channels),
&gt;             nn.ReLU(inplace=True),
&gt;             nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
&gt;             nn.BatchNorm2d(out_channels),
&gt;             nn.ReLU(inplace=True)
&gt;         )
&gt; 
&gt;     def forward(self, x):
&gt;         return self.double_conv(x)
&gt; 
&gt; 
&gt; class Down(nn.Module):
&gt;     """Downscaling with maxpool then double conv"""
&gt; 
&gt;     def __init__(self, in_channels, out_channels):
&gt;         super().__init__()
&gt;         self.maxpool_conv = nn.Sequential(
&gt;             nn.MaxPool2d(2),
&gt;             DoubleConv(in_channels, out_channels)
&gt;         )
&gt; 
&gt;     def forward(self, x):
&gt;         return self.maxpool_conv(x)
&gt; 
&gt; 
&gt; class Up(nn.Module):
&gt;     """Upscaling then double conv"""
&gt; 
&gt;     def __init__(self, in_channels, out_channels, bilinear=True):
&gt;         super().__init__()
&gt; 
&gt;         # if bilinear, use the normal convolutions to reduce the number of channels
&gt;         if bilinear:
&gt;             self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
&gt;             self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
&gt;         else:
&gt;             self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
&gt;             self.conv = DoubleConv(in_channels, out_channels)
&gt; 
&gt; 
&gt;     def forward(self, x1, x2):
&gt;         x1 = self.up(x1)
&gt;         # input is CHW
&gt;         diffY = x2.size()[2] - x1.size()[2]
&gt;         diffX = x2.size()[3] - x1.size()[3]
&gt; 
&gt;         x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
&gt;                         diffY // 2, diffY - diffY // 2])
&gt;         # if you have padding issues, see
&gt;         # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
&gt;         # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
&gt;         x = torch.cat([x2, x1], dim=1)
&gt;         return self.conv(x)
&gt; 
&gt; 
&gt; class OutConv(nn.Module):
&gt;     def __init__(self, in_channels, out_channels):
&gt;         super(OutConv, self).__init__()
&gt;         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
&gt; 
&gt;     def forward(self, x):
&gt;         return self.conv(x)
&gt; 
&gt; class UNet(nn.Module):
&gt;     def __init__(self, n_channels, n_classes, bilinear=True):
&gt;         super(UNet, self).__init__()
&gt;         self.n_channels = n_channels
&gt;         self.n_classes = n_classes
&gt;         self.bilinear = bilinear
&gt; 
&gt;         self.inc = DoubleConv(n_channels, 64)
&gt;         self.down1 = Down(64, 128)
&gt;         self.down2 = Down(128, 256)
&gt;         self.down3 = Down(256, 512)
&gt;         factor = 2 if bilinear else 1
&gt;         self.down4 = Down(512, 1024 // factor)
&gt;         self.up1 = Up(1024, 512 // factor, bilinear)
&gt;         self.up2 = Up(512, 256 // factor, bilinear)
&gt;         self.up3 = Up(256, 128 // factor, bilinear)
&gt;         self.up4 = Up(128, 64, bilinear)
&gt;         self.outc = OutConv(64, n_classes)
&gt; 
&gt;     def forward(self, x):
&gt;         x1 = self.inc(x.float())
&gt;         x2 = self.down1(x1)
&gt;         x3 = self.down2(x2)
&gt;         x4 = self.down3(x3)
&gt;         x5 = self.down4(x4)
&gt;         x = self.up1(x5, x4)
&gt;         x = self.up2(x, x3)
&gt;         x = self.up3(x, x2)
&gt;         x = self.up4(x, x1)
&gt;         logits = self.outc(x)
&gt;        return logits
</code></pre>
<p>Initializing model:</p>
<pre><code class="lang-auto">model = UNet(n_channels=3, n_classes=3, bilinear=False)

optimizer = optim.RMSprop(model.parameters(), lr=0.001, weight_decay=1e-8, momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2)

criterion = nn.CrossEntropyLoss()
</code></pre>
<p>The interesting thing is that when I initialize my model to n_classes = 4,5 and so on, the loop runs just fine, but the outputs are terrible. It’s when I set the n_classes = 3 when it all breaks. Any help would be greatly appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

KFrank: <p>Hi Ben!</p>
<aside class="quote no-group" data-username="benkolber" data-post="1" data-topic="106244" 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/benkolber/48/32373_2.png" class="avatar"> benkolber:</div>
<blockquote>
<p>pixel values of 0,1,2, and 3 (3 classes total, with pixel 0 being the background).</p>
</blockquote>
</aside>
<p>The short story is that you have four classes, not three – “background”<br>
counts as one of your classes.</p>
<aside class="quote no-group">
<blockquote>
<p>When retrieving the output from the model, I have:<br>
output from model = (b_s, 3, 256, 256)<br>
predicted mask = (b_s, 256, 256)</p>
</blockquote>
</aside>
<p>Your model output has shape <code>[b_s, n_classes = 3, 256, 256]</code>.  So,<br>
yes, your model is predicting (only) three classes.</p>
<aside class="quote no-group">
<blockquote>
<p>I use a cross-entropy loss function, but receive the following error:</p>
<pre><code class="lang-auto">IndexError: Target 3 is out of bounds.
</code></pre>
</blockquote>
</aside>
<p><code>CrossEntropyLoss</code> sees that its <code>input</code> (your model output) has<br>
<code>n_classes = 3</code>, so it will require that your <code>target</code> only has values<br>
for three classes.  That is, your <code>target</code> values must be integer class<br>
labels running from <code>[0, n_classes - 1]</code>, i.e., be in <code>(0, 1, 2)</code>.</p>
<p>But your <code>target</code> has values in <code>(0, 1, 2, 3)</code>.  When <code>CrossEntropyLoss</code><br>
hits the “fourth” class value of <code>3</code>, it throws the error you see.</p>
<p>It sounds like your problem really does have four classes, three<br>
foreground classes plus one background class.  So leave your <code>target</code><br>
as it is, but run your model with <code>n_classes = 4</code> so that its output<br>
has shape <code>[b_s, 4, 256, 256]</code>.</p>
<p>Note, you might want to check whether your dataset is <em>unbalanced.</em><br>
It is not unusual to have many more background pixels than foreground.<br>
(I’m not saying you do, but you might.)  If so, you might consider using<br>
<code>CrossEntropyLoss</code>’s <code>weight</code> argument to reweight the classes in your<br>
loss calculation.</p>
<p>Best.</p>
<p>K. Frank</p>

benkolber: <p>Thank you very much! This was very helpful.</p>
<p>Given that I have much more background pixels than classes, I would initialize cross entropy loss as:</p>
<p>criterion = nn.CrossEntropyLoss(weight = tensor([0.3, 1, 1, 1])  to reduce the weight value of the background?</p>
<p>Thanks again! Life saver</p>

KFrank: <p>Hi Ben!</p>
<aside class="quote no-group" data-username="benkolber" data-post="3" data-topic="106244" 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/benkolber/48/32373_2.png" class="avatar"> benkolber:</div>
<blockquote>
<p>Given that I have much more background pixels than classes, I would initialize cross entropy loss as:</p>
<p>criterion = nn.CrossEntropyLoss(weight = tensor([0.3, 1, 1, 1])  to reduce the weight value of the background?</p>
</blockquote>
</aside>
<p>This could be reasonable.  You are weighting each of the foreground<br>
pixels equally, and weighting the background pixels about one third as<br>
much as any one foreground class.</p>
<p>So this would make sense if each of your three foreground classes<br>
appeared about equally often, and <em>each</em> of your foreground classes<br>
<em>individually</em> appeared about one third as often as your background<br>
class.</p>
<p>The typical advice is to weight your classes in about inverse proportion<br>
to their frequency of occurrence.</p>
<p>Having said that, a three-to-one difference in frequency of occurrence<br>
doesn’t strike me as being highly unbalanced.  In such a case I might<br>
expect using class weights to help at the margins, but I doubt doing so<br>
would be a game changer.</p>
<p>Best.</p>
<p>K. Frank</p>