Specify axis for one_hot encoding

MaKaNu: <p>In the following example I have a Tensor with the Dimensions CxWxH with C = 3, W = 2 and H = 2. So I follow the dimensions as an example conv2d is doing. In my later implementation C can be arbitrary value depending on a count of BBoxes in a LabelImage. What I want to achieve is that I reduce the C to 1 by the following set of rules:</p>
<p>If index_C == 0 Any Value above zero should be at its same W and H coordinates in the target tensor.<br>
Else Continue with the next index_C and repeat the previous step until the last index_C.<br>
If still zero values present put them into the target tensor in the same W and H coordinates.</p>
<p>Example</p>
<pre><code class="lang-python">&gt;&gt;&gt;A = torch.tensor(
       [[[0, 1],
         [0, 1]],

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

        [[0, 0],
         [1, 0]]]
)
&gt;&gt;&gt;A.shape
torch.Size([3, 2, 2])
</code></pre>
<p>with a Target which looks like the following:</p>
<pre><code class="lang-python">T = torch.tensor(
    [[0,1],
     [2,1]]
)
</code></pre>
<p>I already achieved this by using the <code>apply_along_axis</code> method and a for loop, but I want to get rid of the loop with broadcasting and vectorization methods.</p>
<p>I achieved so far the following:</p>
<p>With a broadcasting check against 0 I got a True False Tensor:</p>
<pre><code class="lang-python">&gt;&gt;&gt;B = A &gt; 0
&gt;&gt;&gt;B
tensor(
     [[[False,   True],
       [False,   True]],

       [[False, False],
        [ True,  True]],

       [[False, False],
        [ True, False]]]
)
</code></pre>
<p>which I when used to get the <code>argmax</code> index which brings me relative close to my goal:</p>
<pre><code class="lang-python">&gt;&gt;&gt; C = B.long().argmax(axis=0)
&gt;&gt;&gt;C
tensor([[0, 0],
        [1, 0]])
</code></pre>
<p>So what I now only need is a one hot encoding for the indexes. I know that the implementation of Tensorflow assists using the <code>axis</code> argument in their <code>one_hot</code> implementation. So the issue now is the following if I use the <code>one_hot</code> functional implementation I end up with the following:</p>
<pre><code class="lang-python">&gt;&gt;&gt;import torch.nn.functional as F
&gt;&gt;&gt; F.one_hot(C,3)
tensor([[[1, 0, 0],
         [1, 0, 0]],

        [[0, 1, 0],
         [1, 0, 0]]])
</code></pre>
<p>But what I actually needed looks like this:</p>
<pre><code class="lang-python">tensor(
       [[[0, 1],
         [0, 1]],

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

        [[0, 0],
         [0, 0]]]
)
</code></pre>
<p>So it has the same shape as <code>A</code>.</p>
<p>Any Hint how to achieve it without Tensorflow? Our Tensorflow guy implemented it already.</p>

Matias_Vasquez: <p>If your target <code>T</code> matches the last two dimensions of <code>A</code>, then you can simply do something like this</p>
<pre><code class="lang-python">torch.logical_and(A==T, A&gt;0)
</code></pre>
<pre><code class="lang-python"># Output
tensor([[[False,  True],
         [False,  True]],

        [[False, False],
         [ True, False]],

        [[False, False],
         [False, False]]])
</code></pre>
<p>You can then use this as a mask to get the values that you need and where you need them.</p>
<pre><code class="lang-python">B = torch.zeros_like(A) 
B[torch.logical_and(A==T, A&gt;0)] = A[torch.logical_and(A==T, A&gt;0)]
</code></pre>
<pre><code class="lang-python"># Output
tensor([[[0, 1],
         [0, 1]],

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

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

MaKaNu: <p>The Issue is that I want to return the Tensor T. I don’t have Tensor T at the moment. I just represented how T has to look like. So my last example in my post shows the step before the final result. The wanted final result is Tensor T for input Tensor A</p>

Matias_Vasquez: <p>Oh I see.</p>
<p>Then, as you can see, if you permute <code>A</code> to match the output shape of what you did, you can see that the <code>ones</code> correspond to the first value that is taken from <code>A</code>.</p>
<pre><code class="lang-python"># Your answer
&gt;&gt;&gt; F.one_hot(C,3)
tensor([[[1, 0, 0],
         [1, 0, 0]],

        [[0, 1, 0],
         [1, 0, 0]]])

# A permuted
&gt;&gt;&gt; A.permute(1, 2, 0)
tensor([[[0, 0, 0],
         [1, 0, 0]],

        [[0, 2, 1],
         [1, 2, 0]]])
</code></pre>
<p>So, we can permute your answer in the other direction</p>
<pre><code class="lang-python">&gt;&gt;&gt; F.one_hot(C,3).permute(2, 0, 1)
tensor([[[1, 1],
         [0, 1]],

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

        [[0, 0],
         [0, 0]]])
</code></pre>
<p>You can then use this to get <code>T</code>.</p>
<pre><code class="lang-python">mask = F.one_hot(C,3).permute(2, 0, 1)==1
T = torch.zeros_like(A)
T[mask] = A[mask]
T = T.sum(dim=0)
</code></pre>
<pre><code class="lang-python"># Output
tensor([[0, 1],
        [2, 1]])
</code></pre>

Matias_Vasquez: <p>If this assessment of the problem is correct, then you could also do the following[1].</p>
<pre><code class="lang-python">def first_nonzero(x, axis=0):
    nonz = (x &gt; 0)
    return ((nonz.cumsum(axis) == 1) &amp; nonz).max(axis)

A = torch.tensor(
       [[[0, 1],
         [0, 1]],

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

        [[0, 0],
         [1, 0]]]
)

_, inds = first_nonzero(A)
T = A[inds, [[0,0],[1,1]], [[0,1],[0,1]]]
print(T)
</code></pre>
<pre><code class="lang-python"># Output
tensor([[0, 1],
        [2, 1]])
</code></pre>
<p>[1] Function taken from <a href="https://discuss.pytorch.org/t/first-nonzero-index/24769/9">here</a>.</p>

MaKaNu: <p>Wow thanks a lot didn’t thought about <code>permute</code> while struggling with <code>swapaxes</code>. I really like the second attempt and will compare each other. I mark the first one as solved since it fits better the Question title.</p>