How to mask/select specific channel out of segmentation output?

seyeeet: <p>hello everyone,</p>
<p>I am using a segmentation mask, and it gives me <code>output</code> in shape of <code>BxCxHxW</code>, where <code>C</code> is the number of classes.<br>
During my training I only want to optimize for class <code>k</code> where <code>k</code> can be a number  between <code>0</code> to <code>C</code>. So basically I want to only update the model for that specific channel and I dont care about the other channels. in my optimization I use L1 loss to force the <code>k</code> channel to be equal a <code>mask</code>.<br>
Currently, I select my channel via <code>output[:,k]</code> so I will have <code>L1Loss(output[:,k],mask)</code>, where <code>output[:,k] and mask</code> both has shape of <code>BxHxW</code>.<br>
But it seems like something is not working okay. in my output I see that the other channels also get updated , so I assume something is wrong.</p>
<p>I was wondering if someone has a suggestion on how to mask or select a specific channel of  output in a correct way</p>

ptrblck: <p>How are you checking the “updated channels”? Assuming you are using an <code>nn.Conv2d</code> layer as the output layer, each class channel would be created by a separate filter, which would use all input channels by default. So depending what exactly you are checking the updates might be expected.</p>

seyeeet: <p>hmmm that is interesting, what do you mean by <code>How are you checking the “updated channels”</code> ?<br>
I guess I am not, but I am also not sure how should I do it, can you please give me a hint?</p>
<p>here is a psedo code of my code:</p>
<pre><code class="lang-auto">prediction= F.sigmoid(self.deeplabv(inputs)['out']) # self.deeplab is the pytorch deeplab model
batch = prediction.shape[0]//2
output_1 = prediction[:batch]                 
output_2 = prediction[batch:2*batch]
output_1_1 = output_1[selecting_index_1].unsqueeze(1)
output_1_2 = output_2[selecting_index_2].unsqueeze(1)
self.model.eval()
output_1_1_1= self.model(torch.mul(inputs[:batch]         ,output_1_1 ) )
output_1_2_1= self.model(torch.mul(inputs[batch:2*batch]  ,output_1_2 ) )
loss = self.loss(output_1_1_1,mask_1) + self.loss(output_1_2_1,mask_2)
</code></pre>

ptrblck: <p>In your initial post you’ve mentioned:</p>
<aside class="quote no-group" data-username="seyeeet" data-post="1" data-topic="128071">
<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/seyeeet/48/21931_2.png" class="avatar"> seyeeet:</div>
<blockquote>
<p>But it seems like something is not working okay. in my output I see that the other channels also get updated , so I assume something is wrong.</p>
</blockquote>
</aside>
<p>so I wanted to double check what exactly you are checking how and what the expectations are.<br>
The currently posted code snippet also doesn’t show any checks, so I’m unsure what the use case is and what seems to go wrong.</p>

seyeeet: <p>I see, so I was visualizing the output of the final channels, my expectation was that since I only select a specific channel maybe one of them should be similar to the mask, and the other should be different.<br>
however, when I visualize them all of them look very similar.<br>
as an example, if I send a cat image, most of the channels show the mask of the cat, while my expectation was to only see the channel that is correspond to the cat shows the cat. One thought was maybe it is because I am using sigmoid instead of softmax (?) but I was not sure if it still should be like that</p>

RaLo4: <p>I am not sure if I understand you correctly, so I am not sure if this is what you are looking for.</p>
<p>So the output of your model is of shape <code>(B, C, H, W)</code> , correct? And at each location, there are unnormalized probabilities corresponding to the prediction of each class?  Which would be the normal output of, for example <code>torchvision.models.segmentation.deeplabv3_resnet101()['out']</code>.</p>
<p>If so, to then get the maximum prediction of each class, and then use it for a downstream task, you can do <code>output_predictions = output.argmax(0)</code>.</p>
<p>So for example (example with 21 classes):</p>
<pre><code class="lang-auto">with torch.no_grad():
    output = model(input_batch)['out'][0]
output_predictions = output.argmax(0)

# create a color pallette, selecting a color for each class
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
colors = (colors % 255).numpy().astype("uint8")

# plot the semantic segmentation predictions of 21 classes in each color
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
r.putpalette(colors)

import matplotlib.pyplot as plt
plt.imshow(r)
# plt.show()
</code></pre>