How do I get single-channel output when doing semantic segmentation？

username: hao.liu
cooked: <p>Dear:<br>
The problem is that, when i implement the semantic segmetation, the output shape of the semantic segmentation network is [batch_size, num_classes, width, height],it is encoded like one_hot,      but i just want to directly get the output which shape is [batch_size, width, height],in the output each pixel has the value 0,1,2…(indicate the predicted class index), what can i do to get the output’s shape like [batch_size, width, height] (without postprocess like argmax() ).<br>
please help me, thanks all.</p>

username: ptrblck
cooked: <p>In a semantic segmentation use case the output of your model would represent logits and would thus not be one-hour encoded. To get the predicted class label you would use <code>argmax</code>. Since this operation is not differentiable it’s not included in the forward pass of the model.</p>

username: hao.liu
cooked: <p>Hello ptrblck:<br>
Thanks for your reply.</p>
<p>For some reasons, such as deploying neural networks on devices, the .pth fileI will be converted to anonther format like .dlc(made by qualcomm),  what i get is just the result infered by the network, so i have two options: One is to get the output of the shape [batch_size, num_classes, width, height] by the network, and then I need to implement the argmax() function in a language like c++ to get the result of semantic segmentation. The other is that my network output is [batch_size, width, height] so I don’t need to implement the argmax() function myself.</p>
<p>in your reply do you mean that when I use the multi-classes semantic segmentation network, its output shape must be [batch size, num_classes, width, height] ?  and also i need use argmax() function to get the final single-channel result</p>
<p>Thank you.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="hao.liu" data-post="3" data-topic="138000">
<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/h/3be4f8/48.png" class="avatar"> hao.liu:</div>
<blockquote>
<p>in your reply do you mean that when I use the multi-classes semantic segmentation network, its output shape must be [batch size, num_classes, width, height] ? and also i need use argmax() function to get the final single-channel result</p>
</blockquote>
</aside>
<p>Yes, that’s the standard workflow. I’m not familiar with the .dlc format, but in case you can call a custom method of your deployed model, you could also create an e.g. <code>predict</code> method, which could call the <code>forward</code> method with the additional <code>torch.argmax</code> internally. Once deployed, you could then use <code>preds = model.predict(input)</code>.</p>

username: hao.liu
cooked: <p>Hello ptrblck:<br>
Thank you for your kind reply on above question.</p>
<p>I think I understand what you say, you want change the predict fucntion by adding argmax() so that when I run predict function I get the wanted shape.</p>
<p>Except for the method you proposed，can I try this method : I train the model without using argmax() in my forward propagation function( cause it is not differentiable),  when I  use the weight file to do <strong>an inference</strong>, I modify the model( like adding the argmax()  to forward propagation function) so that i can get the result with shape [batch_size, width, height], but I am not sure whether it works.</p>
<p>And also I have another question, if I modify the model( adding argmax() to model) when I implemt inference, i think there may be some probelems when I load the weight file(error may be like :   incompatible keyword etc.)</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="hao.liu" data-post="5" data-topic="138000">
<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/h/3be4f8/48.png" class="avatar"> hao.liu:</div>
<blockquote>
<p>I think I understand what you say, you want change the predict fucntion by adding argmax() so that when I run predict function I get the wanted shape.</p>
<p>Except for the method you proposed，can I try this method : I train the model without using argmax() in my forward propagation function( cause it is not differentiable), when I use the weight file to do <strong>an inference</strong>, I modify the model( like adding the argmax() to forward propagation function) so that i can get the result with shape [batch_size, width, height], but I am not sure whether it works.</p>
</blockquote>
</aside>
<p>Yes, this would be my proposed approach if you can’t add the <code>argmax</code> operation in your deployment setup. Manipulating the <code>forward</code> would work or probably the cleaner approach would be to write a custom <code>predict</code> method which adds the <code>argmax</code> operation to it as described before. You could alsoreturn two outputs in the <code>forward</code>: the logits as well as the class predictions and could use one of them corresponding to training or inference. The “best” approach depends on your workflow.</p>
<aside class="quote no-group" data-username="hao.liu" data-post="5" data-topic="138000">
<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/h/3be4f8/48.png" class="avatar"> hao.liu:</div>
<blockquote>
<p>And also I have another question, if I modify the model( adding argmax() to model) when I implemt inference, i think there may be some probelems when I load the weight file(error may be like : incompatible keyword etc.)</p>
</blockquote>
</aside>
<p>No, you should not see any <code>state_dict</code> mismatches, as <code>torch.argmax</code> is a purely functional call, won’t have any parameters, and will thus not be stored in the <code>state_dict</code>.</p>

username: hao.liu
cooked: <p>Hello ptrblck:</p>
<p>I got what you say. I will try these ways later.</p>
<p>Thank you very much. Have a nice day <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=10" title=":grinning:" class="emoji" alt=":grinning:"></p>