How do i index a pytorch tensor?

1

I have a simple example that is not working. I want to extract a tensor from tr_input using the indices from tr_ind.

    print(np.shape(tr_ind))
    print(np.shape(tr_input))
    
    tr_input1 = tr_input[tr_ind]


torch.Size([150])
torch.Size([1515, 200, 28])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-28-5b55de0850db> in <module>()
      6     print(np.shape(tr_input))
      7 
----> 8     tr_input1 = tr_input[tr_ind,:,:]
      9     #tr_target1 = tr_target[tr_ind]
     10     #val_input = tr_input[val_ind]

~\Anaconda3\envs\dl\lib\site-packages\torch\autograd\variable.py in __getitem__(self, key)
     76                 return IndexSelect.apply(self, 0, key)
     77             # else fall through and raise an error in Index
---> 78         return Index.apply(self, key)
     79 
     80     def __setitem__(self, key, value):

~\Anaconda3\envs\dl\lib\site-packages\torch\autograd\_functions\tensor.py in forward(ctx, i, index)
     87             result = i.index(ctx.index)
     88         else:
---> 89             result = i.index(ctx.index)
     90             ctx.mark_shared_storage((i, result))
     91         return result

TypeError: Performing basic indexing on a tensor and encountered an error indexing dim 0 with an object of type torch.IntTensor. The only supported types are integers, slices, numpy scalars, or if indexing with a torch.LongTensor or torch.ByteTensor only a single Tensor may be passed.
2 Likes
2 
import torch

a = torch.rand(10, 10, 5)
print a[0, :, :]

this works… check the types of tr_ind, that might be the problem

PS: I’m using pytorch v0.4

3

It is a torch int tensor with 150 elements

4

There is also the function index_select which should do exactly what you want

3 Likes
5

Here is a solution if you want to index a tensor in an arbitrary dimension and select a set of tensors from that dimension (an example is say we want to compute some average of the first 3 layers):

# selecting indices arbitrarily i.e. x[*,indicies,*] were * denotes that the rest of the layers are kept the same
# but for only the last 3 layers [T, L] -> [1]
x = torch.randn(5, 4)
# compute average of first 3 layer
L = x.size(1)
indices = torch.tensor(range(L-1))
xx = torch.index_select(x, dim=1, index=indices)
print(f'{x=}')
print(f'{xx=}')
print(xx.size())
x=tensor([[-0.1588, -1.2282, -0.3142, -0.0458],
        [-0.3095, -1.1847,  1.4570, -0.1697],
        [ 0.8444,  0.6528,  0.5663,  2.4564],
        [ 0.3549,  1.8841, -0.6826,  0.0283],
        [ 0.4958,  0.0893,  2.1306, -0.4176]])
xx=tensor([[-0.1588, -1.2282, -0.3142],
        [-0.3095, -1.1847,  1.4570],
        [ 0.8444,  0.6528,  0.5663],
        [ 0.3549,  1.8841, -0.6826],
        [ 0.4958,  0.0893,  2.1306]])
torch.Size([5, 3])

the docs are fantastically clean though: torch.index_select — PyTorch master documentation

actually you can also do:

x.index_select(dim, index)