Mask Shape for Multi Label Image Segmentation

moth: <p>Hello, I was wondering which is the expected shape of a mask for a multilabel image segmentation task. My guess is that each mask’s shape is <code>(channels = num_classes, height, width)</code>, since a segmentation model will output a tensor with as many channels as classes, but i am not entirely sure. Here is an example of a model for a multi-label segmentation task with 5 possible labels.</p>
<pre><code class="lang-auto">model = smp.Unet(
    encoder_name= "efficientnet_b0", # model backbone
    encoder_weights="imagenet", 
    in_channels=3, # 1 grayscale 3 GB
    classes= 5,
    activation=None,
)
</code></pre>
<p>Also, if this is true I guess there is an implicit order between channels and classes. For example, channel1 is always for class1, channel2 is always for class2, …, etc.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> can you confirm this?</p>

KFrank: <p>Hi Moth!</p>
<aside class="quote no-group" data-username="moth" data-post="1" data-topic="172014" data-full="true">
<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/m/aca169/48.png" class="avatar"> moth:</div>
<blockquote>
<p>My guess is that each mask’s shape is <code>(channels = num_classes, height, width)</code></p>
</blockquote>
</aside>
<p>This is basically right with the proviso that the mask tensor you pass to your<br>
loss function should have shape <code>[nBatch, num_classes, height, width]</code>.<br>
Your batch size (<code>nBatch</code>) can be one, but the <code>nBatch</code> dimension is still<br>
required.</p>
<p>Each element of your mask tensor will be a floating point number that is<br>
the probability of that pixel being in the given class.  These values can be<br>
<code>0.0</code> and <code>1.0</code> if you want “hard,” non-probabilistic mask values.</p>
<p>Your predictions – the output of your model – will have the same shape<br>
and consist of <em>logits</em> that correspond to the probability that the given pixel<br>
is in the given class.  These would typically be produced by a final <code>Conv2</code><br>
layer with <code>out_channels = num_classes</code> (and probably a <code>kernel_size</code><br>
of <code>1</code>) and <em>no</em> subsequent <code>sigmoid()</code>.  Your loss function should be<br>
<code>BCEWithLogitsLoss</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>I guess there is an implicit order between channels and classes. For example, channel1 is always for class1, channel2 is always for class2,</p>
</blockquote>
</aside>
<p>I’m not entirely sure what you mean by this.</p>
<p>Your predictions and mask have the same shape.  <code>BCEWithLogitsLoss</code><br>
performs an element-wise binary-cross-entropy computation (and then,<br>
typically a reduction), so the <code>out_channels</code> dimension of your predictions<br>
and the <code>num_classes</code> dimension of your mask do line up.</p>
<p>As aside, these comments apply to the <em>multi-label</em> case where a given<br>
pixel can be assigned to none, some, or all of your class at the same<br>
time.  If instead you are working with what is called the <em>multi-class</em> case,<br>
where a given pixel is in <em>exactly one</em> class at a time, then these comments<br>
don’t directly apply.</p>
<p>Best.</p>
<p>K. Frank</p>

moth: <p>Thanks Frank for the detailed explanation. You confirmed my assumption.</p>
<p>Regarding:</p>
<blockquote>
<p>I guess there is an implicit order between channels and classes. For example, channel1 is always for class1, channel2 is always for class2</p>
</blockquote>
<p>What I meant is that each mask’s channel is always assigned the same category/label. For example, let’s say there are 4 classes: [“cat”, “dog”, “lion”, “wolf”]. We would then have an array/tensor of shape (4, H, W) where the cat mask is always on the first channel of the mask, the dog mask is always on the second channel of the mask, the lion mask if always on the third channel of the mask, the wolf mask is always on the fourth channel of the mask. Essentially, there’s a mapping between categories and channels.</p>
<p>I dont know if I make myself clear. Thanks anyway for your explanation again.</p>
<p>Best wishes</p>