How to create a tensor with a channel for classes

username: WilliManFR
cooked: <p>Hello,<br>
I have a problem. I have done a CNN which gives me an output of the form : (BATCH_SIZE, 1, HEIGHT, WIDTH). However, in order to apply the CrossEntropyLoss function, I need an output of the form, (BATCH_SIZE, NB_CLASS, HEIGHT, WIDTH). How can I get such a matrix from my original output? I currently have three classes, I want to do image segmentation and my output images are obviously float32.</p>
<p>Thank you for your answers <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: KFrank
cooked: <p>Hi William!</p>
<aside class="quote no-group" data-username="WilliManFR" data-post="1" data-topic="125077" data-full="true">
<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/willimanfr/48/38484_2.png" class="avatar"> WilliManFR:</div>
<blockquote>
<p>I have done a CNN which gives me an output of the form : (BATCH_SIZE, 1, HEIGHT, WIDTH).</p>
</blockquote>
</aside>
<p>As an aside, it sounds like your current CNN performs binary (i.e.,<br>
background-foreground) segmentation.</p>
<aside class="quote no-group">
<blockquote>
<p>I need an output of the form, (BATCH_SIZE, NB_CLASS, HEIGHT, WIDTH).</p>
</blockquote>
</aside>
<p>It sounds like your desired use case is multi-class (with <code>NB_CLASS</code><br>
classes) semantic segmentation.</p>
<aside class="quote no-group">
<blockquote>
<p>How can I get such a matrix from my original output?</p>
</blockquote>
</aside>
<p>First, once you’re down to <code>(BATCH_SIZE, 1, HEIGHT, WIDTH)</code>, that<br>
is, a single output channel, you will no longer be able to recover your<br>
desired <code>NB_CLASS</code> channels.  You will have to modify the final bit<br>
of your CNN architecture, and the details will depend on your specific<br>
architecture.</p>
<p>It is likely that your next-to-last layer will produce an “image” with<br>
a number of “feature” channels, and your last layer will recombine<br>
your feature channels into the desired number of “class” channels.</p>
<p>For example, quoting from the original <a href="https://arxiv.org/pdf/1505.04597v1.pdf" rel="noopener nofollow ugc">U-Net</a> paper, “At the final layer<br>
a 1x1 convolution is used to map each 64-component feature vector<br>
to the desired number of classes.”</p>
<p>In the context of pytorch, you might be looking for a final layer of<br>
the form:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.nn.Conv2d (in_channels = 64, out_channels = 1, kernel_size = 1)
</code></pre>
<p>(This would be for the case that your next-to-last layer produces 64<br>
“feature” channels.)</p>
<p>In such a case you would change <code>out_channels</code> to <code>NB_CLASS</code>, that<br>
is, replace your final layer with:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.nn.Conv2d (in_channels = 64, out_channels = NB_CLASS, kernel_size = 1)
</code></pre>
<p>This illustrates the basic idea for classic U-Net – your specific CNN<br>
may differ in detail.</p>
<p>Best.</p>
<p>K. Frank</p>

username: WilliManFR
cooked: <p>Thank you for your answer. It was very clear. I couldn’t get my segmentation to work (it only sends me black images as output) but now my CrossEntropyLoss is working. Thanks a lot!</p>