How to convert argmax result to an one-hot matrix?

BarCodeReader: <p>for example, I have a tensor in shape [N,C,H,W] = [1,3,2,2]<br>
Then I apply softmax and argmax to obtain the index:</p>
<pre><code class="lang-auto"># original tensor
tensor([[[[ 0.4008, -0.6662],
          [-0.4133,  1.3639]],

         [[-0.8354,  0.6317],
          [ 0.3240, -1.1438]],

         [[-0.3452,  1.2110],
          [ 0.6575,  0.9924]]]])

# after softmax
tensor([[[[0.5666, 0.0893],
          [0.1664, 0.5646]],

         [[0.1646, 0.3270],
          [0.3479, 0.0460]],

         [[0.2687, 0.5837],
          [0.4856, 0.3894]]]])

# after argmax on channel dimension
tensor([[[0, 2],
         [2, 0]]])
</code></pre>
<p>then i want to use the index returned by argmax and convert it into a binary matrix:</p>
<pre><code class="lang-auto">tensor([[[[1, 0],
          [0, 1]],  # channel 0

         [[0,0],
          [0,0]],  # channel 1

         [[0,1],
          [1,0]   # channel 2
]]])
</code></pre>
<p>how to do this?</p>

ptrblck: <p>This should work:</p>
<pre><code class="lang-python">x = torch.tensor([[[[ 0.4008, -0.6662],
                    [-0.4133,  1.3639]],

                   [[-0.8354,  0.6317],
                    [ 0.3240, -1.1438]],

                   [[-0.3452,  1.2110],
                    [ 0.6575,  0.9924]]]])

pred = torch.argmax(x, dim=1)
print(pred)
&gt; tensor([[[0, 2],
           [2, 0]]])

out = torch.zeros_like(x).scatter_(1, pred.unsqueeze(1), 1.)
print(out)
&gt; tensor([[[[1., 0.],
            [0., 1.]],

           [[0., 0.],
            [0., 0.]],

           [[0., 1.],
            [1., 0.]]]])
</code></pre>

BarCodeReader: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="125508">
<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><code>torch.zeros_like(x).scatter_(1, pred.unsqueeze(1), 1.)</code></p>
</blockquote>
</aside>
<p>Much appreciate! it works</p>

KFrank: <p>Hi Yun and <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="125508" 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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">out = torch.zeros_like(x).scatter_(1, pred.unsqueeze(1), 1.)
</code></pre>
</blockquote>
</aside>
<p>As an alternative, you could use <code>torch.nn.functional.one_hot()</code>,<br>
but would now have to <code>permute()</code> the dimensions of the result to get<br>
them back in the order you want them:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt;
&gt;&gt;&gt; t = torch.tensor ([
...     [[[ 0.4008, -0.6662],
...     [-0.4133,  1.3639]],
...
...     [[-0.8354,  0.6317],
...     [ 0.3240, -1.1438]],
...
...     [[-0.3452,  1.2110],
...     [ 0.6575,  0.9924]]]
... ])
&gt;&gt;&gt;
&gt;&gt;&gt; ind = t.argmax (dim = 1)
&gt;&gt;&gt;
&gt;&gt;&gt; ind
tensor([[[0, 2],
         [2, 0]]])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.nn.functional.one_hot (ind).permute (0, 3, 1, 2)
tensor([[[[1, 0],
          [0, 1]],

         [[0, 0],
          [0, 0]],

         [[0, 1],
          [1, 0]]]])
</code></pre>
<p>This is arguably a bit more readable, as with “<code>one_hot()</code>” you say<br>
what you mean (at the cost of the <code>permute()</code>).</p>
<p>(Note that, as <a class="mention" href="/u/ptrblck">@ptrblck</a> illustrated, you do not need to use <code>softmax()</code>.<br>
Because it doesn’t change the <em>order</em> of the values, leaving it out<br>
doesn’t change the result of <code>argmax()</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>