Specify axis for one_hot encoding

vision
1

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:

If index_C == 0 Any Value above zero should be at its same W and H coordinates in the target tensor.
Else Continue with the next index_C and repeat the previous step until the last index_C.
If still zero values present put them into the target tensor in the same W and H coordinates.

Example

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

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

        [[0, 0],
         [1, 0]]]
)
>>>A.shape
torch.Size([3, 2, 2])

with a Target which looks like the following:

T = torch.tensor(
    [[0,1],
     [2,1]]
)

I already achieved this by using the apply_along_axis method and a for loop, but I want to get rid of the loop with broadcasting and vectorization methods.

I achieved so far the following:

With a broadcasting check against 0 I got a True False Tensor:

>>>B = A > 0
>>>B
tensor(
     [[[False,   True],
       [False,   True]],

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

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

which I when used to get the argmax index which brings me relative close to my goal:

>>> C = B.long().argmax(axis=0)
>>>C
tensor([[0, 0],
        [1, 0]])

So what I now only need is a one hot encoding for the indexes. I know that the implementation of Tensorflow assists using the axis argument in their one_hot implementation. So the issue now is the following if I use the one_hot functional implementation I end up with the following:

>>>import torch.nn.functional as F
>>> F.one_hot(C,3)
tensor([[[1, 0, 0],
         [1, 0, 0]],

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

But what I actually needed looks like this:

tensor(
       [[[0, 1],
         [0, 1]],

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

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

So it has the same shape as A.

Any Hint how to achieve it without Tensorflow? Our Tensorflow guy implemented it already.

2

If your target T matches the last two dimensions of A, then you can simply do something like this

torch.logical_and(A==T, A>0)

# Output
tensor([[[False,  True],
         [False,  True]],

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

        [[False, False],
         [False, False]]])

You can then use this as a mask to get the values that you need and where you need them.

B = torch.zeros_like(A) 
B[torch.logical_and(A==T, A>0)] = A[torch.logical_and(A==T, A>0)]

# Output
tensor([[[0, 1],
         [0, 1]],

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

        [[0, 0],
         [0, 0]]])
3

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

4

Oh I see.

Then, as you can see, if you permute A to match the output shape of what you did, you can see that the ones correspond to the first value that is taken from A.

# Your answer
>>> F.one_hot(C,3)
tensor([[[1, 0, 0],
         [1, 0, 0]],

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

# A permuted
>>> A.permute(1, 2, 0)
tensor([[[0, 0, 0],
         [1, 0, 0]],

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

So, we can permute your answer in the other direction

>>> F.one_hot(C,3).permute(2, 0, 1)
tensor([[[1, 1],
         [0, 1]],

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

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

You can then use this to get T.

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)

# Output
tensor([[0, 1],
        [2, 1]])
5

If this assessment of the problem is correct, then you could also do the following[1].

def first_nonzero(x, axis=0):
    nonz = (x > 0)
    return ((nonz.cumsum(axis) == 1) & 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)

# Output
tensor([[0, 1],
        [2, 1]])

[1] Function taken from here.

1 Like
6

Wow thanks a lot didn’t thought about permute while struggling with swapaxes. 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.

1 Like