Confuse about torch.tensor's indexing

1

Can anyone tell me, What’s the difference between mask_a and mask_b?
thanks!

import torch
torch.manual_seed(1)
data = torch.rand(3,3)
print('data',data,'\n')

mask_a = a[:,0]>0.5
mask_a = mask_a.unsqueeze(-1)
mask_a = mask_a.expand_as(data)
print('mask_a',data[mask_a],'\n')

mask_b = torch.tensor([1,1,0])
mask_b = mask_b.unsqueeze(-1)
mask_b = mask_b.expand_as(data)
print('mask_b',data[mask_b],'\n')

output

data tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999],
        [0.3971, 0.7544, 0.5695]]) 

mask_a tensor([0.7576, 0.2793, 0.4031, 0.7347, 0.0293, 0.7999]) 

mask_b tensor([[[0.7347, 0.0293, 0.7999],
         [0.7347, 0.0293, 0.7999],
         [0.7347, 0.0293, 0.7999]],
        [[0.7347, 0.0293, 0.7999],
         [0.7347, 0.0293, 0.7999],
         [0.7347, 0.0293, 0.7999]],
        [[0.7576, 0.2793, 0.4031],
         [0.7576, 0.2793, 0.4031],
         [0.7576, 0.2793, 0.4031]]])
2

The first example uses “boolean array indexing” while the second uses plain indexing as described in the numpy advanced indexing docs.
If you call print('mask_b', data[mask_b.byte()]), you’ll get the same result.

1 Like
3

thanks man. help a lot !