Multiclass semantic segmentation with RGB masks steps clarification

kristina: <p>Hi</p>
<p>I am really confused about the multi-class semantic segmentation steps</p>
<p><strong>I have</strong></p>
<ul>
<li>1000 Images</li>
<li>6 classes of ground truth RGB masks as PNGs</li>
</ul>
<p>I am using <code>nn.CrossEntropyLoss</code> as my loss function.<br>
Evaluation metric is IoU</p>
<p><strong>My first approach was to</strong></p>
<ol>
<li>Convert original ground truth masks to <code>one_hot_encoded</code> based on pixel colour</li>
<li>Convert the images and masks to <code>compose.ToTensor()</code></li>
<li>Train and use IOU metrics for evaluations</li>
</ol>
<p>Then I learnt that crossEntropy suppose to take categorical indexed values 1, 2, 3, 4, 5, 6</p>
<p><strong>My current approach</strong></p>
<ol>
<li>Convert original ground truth masks to Class categories. Each pixel belong to specific class</li>
<li>Convert the class categorised masks to tensor using <code>torch.tensor()</code></li>
<li>Convert the images to <code>compose.ToTensor()</code></li>
<li>NB: Note the difference functions used to convert to tensor.</li>
</ol>
<p>Issues &amp; questions</p>
<ol>
<li>However, I seem to loose class dimension? why is that?</li>
<li>Metric to use - I see that IoU take pixel wise inputs? do I need to somehow convert the categorical predicted class &amp; categorical ground truth masks and convert it back to its original colour code before evaluate?</li>
</ol>

ptrblck: <aside class="quote no-group" data-username="kristina" data-post="1" data-topic="184951">
<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/k/2acd7d/48.png" class="avatar"> kristina:</div>
<blockquote>
<p>However, I seem to loose class dimension? why is that?</p>
</blockquote>
</aside>
<p>I assume you are wondering about the shape of the target which should be <code>[batch_size, height, width]</code> (note the missing “class” dimension) containing class indices in the range <code>[0, nb_classes-1]</code>?<br>
If so, then note that the class dimension would be useless since you are already using indices to represent the class.</p>

kristina: <p>So do I need to do something like this?</p>
<pre><code class="lang-auto">for x, y in iterator:
  x, y = x.to(self.device), y.to(self.device)
  y = torch.unsqueeze(y, x.shape[1])` not sure if this is correct?
</code></pre>
<p>Adopted from <a href="https://github.com/qubvel/segmentation_models.pytorch/blob/e5d3db20e9c2ddb76f88642409e527239943c983/segmentation_models_pytorch/utils/train.py#L48-L52" rel="noopener nofollow ugc">https://github.com/qubvel/segmentation_models.pytorch/blob/e5d3db20e9c2ddb76f88642409e527239943c983/segmentation_models_pytorch/utils/train.py#L48-L52</a></p>
<pre><code class="lang-auto">class TrainEpoch(Epoch):

    def __init__(self, model, loss, metrics, optimizer, logger, device='cpu', verbose=True, writer=None ):
        super().__init__(
            model=model,
            loss=loss,
            logger=logger,
            metrics=metrics,
            stage_name='train',
            device=device,
            verbose=verbose,
            writer=None,
        )
        self.writer = writer
        self.optimizer = optimizer
        self.logger = logger
        self.log_on_start = True

    def on_epoch_start(self):
        self.model.train()

    def batch_update(self, x, y):
        print('Shape',  x.shape, y.shape)

        self.optimizer.zero_grad()
        prediction = self.model.forward(x)

        if isinstance(prediction, dict):
            prediction = prediction['out']

            if self.log_on_start:
                Warning("prediction is a dictionary, using 'out' key")
                self.logger.warning("prediction is a dictionary, using 'out' key")
                self.log_on_start = False
        print(prediction.shape, y.shape)
        loss = self.loss(prediction, y)
        loss.backward()
        self.optimizer.step()
        return loss, prediction
</code></pre>
<p>I get some errors like this</p>
<pre><code class="lang-auto">           for x, y in iterator:
                x, y = x.to(self.device), y.to(self.device)
                print('Shape iterator',  x.shape, y.shape, torch.unsqueeze(y, dim=x.shape[1]).shape, x.shape[1])

                # Add number of channels to y
                        # Get the number of channels from tensor x
                num_channels = x.size(1)

                # If y lacks the channel dimension, add it with the same number of channels as x
                y = y.unsqueeze(1).expand(-1, num_channels, -1, -1)
</code></pre>
<pre><code class="lang-auto">
Shape iterator torch.Size([2, 3, 512, 512]) torch.Size([2, 512, 512]) torch.Size([2, 512, 512, 1]) 3
new shape torch.Size([2, 3, 512, 512]) torch.Size([2, 3, 512, 512])
Shape Train torch.Size([2, 3, 512, 512]) torch.Size([2, 3, 512, 512])
prediction shape torch.Size([2, 8, 512, 512]) torch.Size([2, 3, 512, 512])

train Epoch 0:   0%|          | 0/35 [00:04&lt;?, ?it/s]
Traceback (most recent call last):
  File "/kristina/dev/training/UNet/train_smp.py", line 96, in &lt;module&gt;
    train_logs = train_epoch.run(_train_loader, epoch=i)
  File "/kristina/dev/training/UNet/../smp/utils/train.py", line 65, in run
    loss, y_pred = self.batch_update(x, y)
  File "/kristina/dev//-training/UNet/../smp/utils/train.py", line 139, in batch_update
    loss = self.loss(prediction, y)
  File "/miniconda3/envs/conda_env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/miniconda3/envs/conda_env/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/kristina/miniconda3/envs/conda_env/lib/python3.10/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of size: : [2, 3, 512, 512]
</code></pre>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <aside class="quote no-group" data-username="kristina" data-post="3" data-topic="184951">
<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/k/2acd7d/48.png" class="avatar"> kristina:</div>
<blockquote>
<p><code>RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of size: : [2, 3, 512, 512]</code></p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="184951">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>shape of the target which should be <code>[batch_size, height, width]</code> (note the missing “class” dimension) containing class indices in the range <code>[0, nb_classes-1]</code>?</p>
</blockquote>
</aside>
<p>The error message prints the invalid target shape so check my previous post to see which shape is expected.</p>

kristina: <p>I see! I think I misunderstood what I have to do. I guess I should not have out_classes in my output? so for</p>
<pre><code class="lang-auto"># Final Convolution UNET
self.conv_last = nn.Conv2d(64, out_classes, kernel_size=1)
</code></pre>
<pre><code class="lang-auto">       # Apply softmax activation along the 'out_classes' dimension 
       prediction_probs = F.softmax(prediction, dim=1)

        # Take argmax along the 'out_classes' dimension to get the final class predictions
        predicted_classes = torch.argmax(prediction_probs, dim=1)

        loss = self.loss(predicted_classes, y.float())
        print('loss', loss)
        loss.backward()
</code></pre>
<p>Then I get</p>
<pre><code class="lang-auto">RuntimeError: "host_softmax" not implemented for 'Long'
srun: error: gcn4: task 0: Exited with exit code 1
srun: Terminating StepId=3138493.0
</code></pre>
<p>Thanks for the help in advance <img src="https://discuss.pytorch.org/images/emoji/apple/bowing_woman.png?v=12" title=":bowing_woman:" class="emoji" alt=":bowing_woman:" loading="lazy" width="20" height="20"></p>

kristina: <p>and this code<br>
<code> loss = self.loss(predicted_classes.float(), y.float())</code></p>
<p>gives me this error</p>
<pre><code class="lang-auto">    loss, y_pred = self.batch_update(x, y)
  File "/dev/UNet/../smp/utils/train.py", line 141, in batch_update
    loss.backward()
  File "/ miniconda3/envs/conda_env/lib/python3.10/site-packages/torch/_tensor.py", line 487, in 
backward
    torch.autograd.backward(
  File "/ miniconda3/envs/conda_env/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
srun: error: gcn4: task 0: Exited with exit code
</code></pre>

ptrblck: <p>No, the model output is expected to contain raw logits (so remove the <code>F.softmax</code>) in the shape <code>[batch_size, nb_classes, height, width]</code> while the target should contain class indices in the range <code>[0, nb_classes-1]</code> and the shape <code>[batch_size, height, width]</code>. The previous error pointed to the target, not the model output.</p>