Unfolding tensor based on binary map

vision
1

I have a 4d tensor that I need to perform some operations along the fourth dimension.

I have a 2d binary map for the first two dimensions, where I need to perform the operation only if it’s true.

I found that I can do this to get the sub tensor that I need to perform the operations on:

sub_patches=patches[binary_map, :, :]

which gives the desired time speed up.

but this combines the first two dimension, I need to put back where the values were.

Does anybody know how to do that?

Thanks

2

Can you just assign back the sub tensor to the original at the desired index?

t = torch.arange(0, 36).view(2, 2, 3, 3)
binary_map = (torch.arange(0, 4) % 2 == 0).view(2, 2)
sub_t = t[binary_map, :, :]
sub_t = sub_t * 0
t[binary_map, :, :] = sub_t  # just assigning it back
print(t)

Output:
tensor([[[[ 0,  0,  0],
          [ 0,  0,  0],
          [ 0,  0,  0]],

         [[ 9, 10, 11],
          [12, 13, 14],
          [15, 16, 17]]],


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

         [[27, 28, 29],
          [30, 31, 32],
          [33, 34, 35]]]])
3

thanks. I tried this and thought things were put back in different order. It turns out it’s because I was calculating median and the result I was comparing against took the smaller number when there’s even number, while pytorch took the average.