Segmentation Network Loss issues

username: Alex_Ge
cooked: <p>Hello,</p>
<p>I’ve read quite a few relevant topics here on <code>discuss.pytorch.org</code> such as:</p>
<ul>
<li><a href="https://discuss.pytorch.org/t/loss-function-for-segmentation-models/32129/3" class="inline-onebox">Loss function for segmentation models</a></li>
<li><a href="https://discuss.pytorch.org/t/convert-pixel-wise-class-tensor-to-image-segmentation/1268" class="inline-onebox">Convert pixel wise class tensor to image segmentation</a></li>
<li><a href="https://discuss.pytorch.org/t/fcn-implementation-loss-function/8435/4" class="inline-onebox">FCN Implementation : Loss Function</a></li>
</ul>
<p>I’ve tried with <code>CrossEntropyLoss</code> but it comes with problems I don’t know how to easily overcome.<br>
So I’m now trying to use <code>nn.NLLLoss</code> with pytorch <code>1.3</code> after running the network logits through <code>torch.nn.functional.log_softmax</code></p>
<p>The way the dataset is organised, is that each input has a 7-channel tensor, where each channel is a class, and then the dims (224x224) are the pixels for that class.<br>
So because I can have multiple classes in an image, I want them all indexed, across their respective channels.</p>
<p>I’ve followed the instructions from the above topics, but I keep getting the same error below:</p>
<pre><code class="lang-auto">logits: torch.Size([32, 7, 224, 224])
target: torch.Size([32, 7, 224, 224])
Traceback (most recent call last):
  File "train_segnet.py", line 183, in &lt;module&gt;
    loss   = criterion(logits, labels.cuda().long())
  File "venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "venv/lib/python3.6/site-packages/torch/nn/modules/loss.py", line 204, in forward
    return F.nll_loss(input, target, weight=self.weight, ignore_index=self.ignore_index, reduction=self.reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 1840, in nll_loss
    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size: : [32, 7, 224, 224]
</code></pre>
<p>I thought that <code>NLLoss</code> is capable of dealing with multiple dimension tensors?<br>
Should I squeeze them all across a single axis? How do I solve this issue so that the network<br>
can learn to produce those tensors as output?</p>

username: ptrblck
cooked: <p>Based on your description I assume you are working on a multi-class segmentation use case, i.e. each pixel belongs to one of the 7 classes.<br>
In that case, the target should not contain the channel dimension (number of classes), but instead have the shape <code>[batch_size, height, width]</code> and contain values in the range <code>[0, nb_classes]</code>.</p>
<p>Here is a small dummy code using your shapes:</p>
<pre><code class="lang-python">nb_classes = 7
batch_size, height, width = 32, 224, 224

output = torch.randn(batch_size, nb_classes, height, width, requires_grad=True)
target = torch.randint(0, nb_classes, (batch_size, height, width))

criterion = nn.CrossEntropyLoss()

loss = criterion(output, target)
loss.backward()
</code></pre>

username: Alex_Ge
cooked: <p>What’s the advantage of doing this? This is the way I originally tried to implement it, but I couldn’t work out the accuracy or loss implementation. I’ll try doing it this way and see what happens <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Also, would you recommend <code>CrossEntropyLoss</code>, <code>NLLloss</code> or some other function for this sort of thing, and why one over the other?</p>

username: Alex_Ge
cooked: <p>So I tried what was suggested, and I’m still quite confused, since now I am getting the following error:</p>
<pre><code class="lang-auto">logit shape torch.Size([32, 224, 224])
label shape torch.Size([32, 224, 224])
Traceback (most recent call last):
  File "train_segnet.py", line 185, in &lt;module&gt;
    loss   = criterion(logits, labels.cuda().long())
  File "venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "venv/lib/python3.6/site-packages/torch/nn/modules/loss.py", line 916, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 2009, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 1848, in nll_loss
    out_size, target.size()))
ValueError: Expected target size (32, 224), got torch.Size([32, 224, 224])
</code></pre>
<p>However:</p>
<ol>
<li>the network seems to output values in a different range (I am using <code>DeepLabV3('resnet18', classes=1)</code> which is from this repository here: <a href="https://github.com/qubvel/segmentation_models.pytorch" rel="nofollow noopener">https://github.com/qubvel/segmentation_models.pytorch</a>
</li>
<li>that was the reason I got everything placed in different channels (the repo uses that approach)</li>
<li>I’ve placed my data as you suggested, however that doesn’t seem to work for me right now.</li>
</ol>
<p>One of the errors I am currently getting is:</p>
<p><code>/pytorch/aten/src/THCUNN/SpatialClassNLLCriterion.cu:104: void cunn_SpatialClassNLLCriterion_updateOutput_kernel(T *, T *, T *, long *, T *, int, int, int, int, int, long) [with T = float, AccumT = float]: block: [7,0,0], thread: [991,0,0] Assertion</code>t &gt;= 0 &amp;&amp; t &lt; n_classes<code>failed</code></p>
<p>Which makes sense because when I look at the network output/logits they are:</p>
<pre><code class="lang-auto">[ 0.3511,  0.3759,  0.4006,  ...,  0.2662,  0.2667,  0.2672],
          [ 0.3513,  0.3746,  0.3980,  ...,  0.3169,  0.3135,  0.3100],
          [ 0.3514,  0.3734,  0.3954,  ...,  0.3677,  0.3602,  0.3528],
          ...,
          [ 0.2786,  0.2707,  0.2628,  ...,  0.0179, -0.0218, -0.0616],
          [ 0.2405,  0.2318,  0.2231,  ..., -0.0140, -0.0691, -0.1243],
          [ 0.2024,  0.1928,  0.1833,  ..., -0.0458, -0.1164, -0.1870]]],


        [[[-0.1596, -0.0971, -0.0346,  ...,  0.2567,  0.2747,  0.2928],
          [-0.0786, -0.0282,  0.0221,  ...,  0.2614,  0.2682,  0.2750],
          [ 0.0025,  0.0406,  0.0787,  ...,  0.2662,  0.2617,  0.2572],
          ...,
          [ 0.2729,  0.2585,  0.2440,  ..., -0.1774, -0.2158, -0.2543],
          [ 0.2877,  0.2605,  0.2333,  ..., -0.2256, -0.2790, -0.3325],
          [ 0.3024,  0.2625,  0.2225,  ..., -0.2739, -0.3423, -0.4107]]],


        [[[ 0.0162,  0.0238,  0.0313,  ...,  0.0944,  0.0774,  0.0605],
          [ 0.0553,  0.0645,  0.0736,  ...,  0.1529,  0.1334,  0.1138],
          [ 0.0943,  0.1052,  0.1160,  ...,  0.2115,  0.1893,  0.1671],
          ...,
          [-0.0165,  0.0217,  0.0599,  ..., -0.0660, -0.0712, -0.0764],
          [-0.0538, -0.0219,  0.0100,  ..., -0.0507, -0.0527, -0.0547],
          [-0.0912, -0.0655, -0.0399,  ..., -0.0355, -0.0342, -0.0329]]],


        ...,
</code></pre>
<p>So, is there some transformation I should apply to the logits?<br>
Or should I go back to dealing with 7-channel output, and then translate those, to a 1-dim tensor which can be used to calculate losses, in the shape you have suggested?<br>
I can obviously go from 7-channel per class, to 1-channel with class index values, and then find a way to translate those so that the losses can be calculated, assuming that will work as per your example.</p>

username: ptrblck
cooked: <p>Your logit output shape is missing the class dimension.<br>
In my code snippet I’m creating the logits as <code>[batch_size, nb_classes, height, width]</code> and the target es <code>[batch_size, height, width]</code>. If you stick to these shapes, it should work.</p>
<aside class="quote no-group" data-username="Alex_Ge" data-post="3" data-topic="73797">
<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/alex_ge/48/17912_2.png" class="avatar"> Alex_Ge:</div>
<blockquote>
<p>Also, would you recommend <code>CrossEntropyLoss</code> , <code>NLLloss</code> or some other function for this sort of thing, and why one over the other?</p>
</blockquote>
</aside>
<p><code>nn.CrossentropyLoss</code> expects logits and uses <code>F.log_softmax</code> + <code>nn.NLLLoss</code> internally, so these approaches will yield the same result.</p>

username: Alex_Ge
cooked: <p>I see now why this has been so confusing. I’ve run your example and realised what’s going on.<br>
The actual logits are of the shape of <code>[batch, class_channel, height, width]</code> which is what I’ve had all along. The labels/targets/ideals are in the shape of <code>[batch, height, width]</code>. That is the confusing part, and I didn’t really get why it’s like this, until I saw what was happening at the logits:</p>
<ul>
<li>if they didn’t have a separate channel for the <code>class</code> or <code>label</code> I suspect there is no way to transform their raw/logit value to a class representation? Which is the problem I’ve run into twice so far.<br>
So I transformed the labels into a 1-dim tensor, where each pixel represents a class index:</li>
</ul>
<p>So I wrote this:</p>
<pre><code class="lang-auto">def segnet_label_translate(labels, num_classes):
    labels = labels.cuda().long()
    cuda0  = torch.device('cuda:0')
    b_size = labels.shape[0]
    tensor = torch.zeros([b_size, 224, 224], dtype=torch.long, device=cuda0)
    zeros  = torch.zeros([224, 224], dtype=torch.long, device=cuda0)

    for batch in range(b_size):
        for i in range(num_classes):
            ideal    = labels[batch][i]
            label    = i  + 1
            values   = torch.full([224, 224], label, dtype=torch.long, device=cuda0)
            filtered = torch.where((ideal != 0), values, zeros)
            tensor[batch] = filtered

    return tensor
</code></pre>
<p>I suspect there are easier of faster ways to write the above^.<br>
That seems to work when I remove the line which does <code>label = i + 1</code> otherwise I get the assertion triggered. However, what is unclear is why would I use class zero.<br>
Then I remembered reading somewhere that class <code>0</code> in segmentation is always the background or some other <strong>less useful</strong> class?</p>
<p>The above works if I don’t do <code>label = i + 1</code> but for me class <code>0</code> is something valid that I look for, so I suppose I’ll have to account for that.</p>
<p>Also, should I be applying a <code>softmax</code> or any other transformation to the logits, or is there no need for it?</p>

username: ptrblck
cooked: <p>I’m not completely sure, what your translate method is doing.<br>
How did you define the target at the moment?</p>
<aside class="quote no-group" data-username="Alex_Ge" data-post="6" data-topic="73797">
<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/alex_ge/48/17912_2.png" class="avatar"> Alex_Ge:</div>
<blockquote>
<p>However, what is unclear is why would I use class zero.<br>
Then I remembered reading somewhere that class <code>0</code> in segmentation is always the background or some other <strong>less useful</strong> class?</p>
</blockquote>
</aside>
<p>Class0 has the same importance as any other class. It is often used as the background, but you can assign whatever class you want to it.<br>
Note that Python uses 0-indexing, which makes class0 just the first class.</p>
<aside class="quote no-group" data-username="Alex_Ge" data-post="6" data-topic="73797">
<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/alex_ge/48/17912_2.png" class="avatar"> Alex_Ge:</div>
<blockquote>
<p>Also, should I be applying a <code>softmax</code> or any other transformation to the logits, or is there no need for it?</p>
</blockquote>
</aside>
<p>No, as said above, <code>nn.CrossentropyLoss</code> expects logits and uses <code>F.log_softmax</code> + <code>nn.NLLLoss</code> internally, so you should not add a softmax layer.</p>

username: Alex_Ge
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> It takes the current  <code>[batch_size, nb_classes, height, width]</code> target which is spread across 7 channels (8 now in order to add background) and then squeezes the <code>nb_classes</code> dimension from 8 to 1, since my pixels/classes are never overlapping (I’m dealing with documents).</p>
<p>Hence it iterates the  <code>[batch_size, nb_classes, height, width]</code> and transforms it to <code>[batch_size, height, width]</code> instead, whereas before each channel had a <strong>binary</strong> value for each pixel, it now has the equivalent <code>long</code> integer representing the class, as you suggested.</p>
<p>I’m pretty sure someone else (you most likely) can do the above in a simpler and more elegant way.<br>
PS: <code>b_size</code> is batch size, and <code>labels[batch]</code> is the <code>[nb_classes, height, width]</code> tensor.</p>

username: ptrblck
cooked: <p>Thanks for the explanation.<br>
If each channel is one-hot encoded, you could try to use <code>target = torch.argmax(target, dim=1)</code> instead.</p>

username: Alex_Ge
cooked: <p>Also, I forgot to ask. When deployed, how can I make sense of the logits in order to translate them into the pixel coordinates of a specific class?</p>

username: ptrblck
cooked: <p>You get the predicted class for each pixel also via: <code>preds = torch.argmax(output, dim=1)</code>.<br>
This will return the indices for the maximal values in dim1, which corresponds to the class dimension.</p>

username: Alex_Ge
cooked: <p>That’s interesting. This seems to return only one of the labels, not the others which I know are present (and the network has been trained). What could cause this?</p>
<p>I’m uploading an example of what I mean,</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d50c51a95054b82989e9945692ad60d85b962214.jpeg" alt="3ed9367921a6a734676def01597407ab2e47aa47dd5376c07315a0f9be09e087c3f9704d8ce0c5d2546e30167747657d04527176911b5892836251c87c95f660" data-base62-sha1="uoI5RkE5jUj3p5aVzcbzqxzzhk0" width="158" height="224"></p>
<p>In this picture (it is a human-readable mask) I’m only getting the green region (class 7) but not the other two regions (class 2 and class 3).</p>
<p>This is trained with the latest samples above, and with a pixel-wise accuracy of 99%.</p>

username: ptrblck
cooked: <p>How are you calculating the pixel-wise accuracy?<br>
You would have to get the predicted classes somehow to calculate it, so you could also use it to visualize the prediction.</p>

username: Alex_Ge
cooked: <p>I am not sure that using Cross Entropy does what I want it to; I suspect it trains it to maximise the class index, however in a single image I can have multiple classes present at various pixel locations. So far the network seems to learn to detect only the pixels of the class with the highest index?</p>
<p>Pixel wise accuracy is done by the following snippet (I had assumed softmax was needed), as such the code below won’t work without it, that is because each channel-per-class had a binary value.</p>
<pre><code class="lang-auto">           input = images.cuda(non_blocking=True).float()
            labels = labels.cuda(non_blocking=True).long()
            outputs = model(input)

            b_size = labels.shape[0]

            # iterate labels in dim 0, e.g., batch size
            for i in range(b_size):
                #
                #   squeeze and transform the [7,224,224] to
                #   a single dim [351232] of pixel values 
                #   all concat across one dim for all channels
                #
                label   = labels[i].view(1,-1).squeeze()
                actual  = f(outputs[i].view(1, -1).squeeze())
                l_size  = labels[i].shape[0]
                #
                #   pixel-wise comparison (note this requires softmaxed/min-maxed values)
                #
                actual  = (actual &gt; 0.9).float() * 1
                same    = (actual == label).sum()
                total   = l_size * 224 * 224
                match   = same.item() / total
                #
                # 50176 is the 224 x 224 so we need to 
                # divide and normalise the correct values now
                #
                correct.append(match)
                total   += 1
</code></pre>
<p>I now suspect that what you had originally suggested (a target with single dim) would make more sense as it would allow me to argmax and calculate accuracy based on that.</p>

username: Alex_Ge
cooked: <p>So I’ve reverted to using targets which are in shape of <code>[batch, height, width]</code> and their values are the class index (so no more channel per class).</p>
<p>Which AFAIK is what you have suggested. The loss function seems to work without hacking into any type of tensor or channels, simply with <code>torch.argmax</code>, however now my accuracy doesn’t work on a pixel-wise level, so I’m trying to implement it this way:</p>
<pre><code class="lang-auto">            outputs  = model(input)
            preds    = torch.argmax(outputs, dim=1)
            matching += (preds == labels).sum().item()
            correct  += (matching / 224 * 224)

            b_size  = labels.shape[0]
            total   += b_size
</code></pre>
<p>This seems to give a lower Top-1 score than before, but not by much (98% vs 99%).</p>
<p>So I am testing it again with the deployment platform, and after running the following:</p>
<pre><code class="lang-auto">        t_in   = self.tf(image)
        t_in   = torch.unsqueeze(t_in, 0)
        logits = self.net(t_in)
        torch.set_printoptions(profile="full")
        torch.set_printoptions(precision=5)
        #logits = logits.squeeze()
        print(logits.shape)
        print(logits)
        preds  = torch.argmax(logits, dim=1)
        print(preds.shape)
        print(preds)
        return self.extract(image, logits)
</code></pre>
<p>I noticed that the <code>pred</code> now contains the pixel class index correctly for all the classes present in the document.</p>
<p>So, <strong>THANK YOU</strong>, I owe you a few rounds of beer, your help <strong>is greatly appreciated</strong>!!!</p>

username: ptrblck
cooked: <p>Haha, good to hear it’s working now! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>