For example, I got a tensor:
tensor = torch.rand(12, 512, 768)
And I got an index list, say it is:
[0,2,3,400,5,32,7,8,321,107,100,511]
I wish to select 1 element out of 512 elements on dimension 2 given the index list. And then the tensor's size would become (12, 1, 768).
Is there a way to do it?
There is also a way just using PyTorch and avoiding the loop using indexing and torch.split:
tensor = torch.rand(12, 512, 768)
# create tensor with idx
idx_list = [0,2,3,400,5,32,7,8,321,107,100,511]
# convert list to tensor
idx_tensor = torch.tensor(idx_list)
# indexing and splitting
list_of_tensors = tensor[:, idx_tensor, :].split(1, dim=1)
When you call tensor[:, idx_tensor, :] you will get a tensor of shape:
(12, len_of_idx_list, 768).
Where the second dimension depends on your number of indices.
Using torch.split this tensor is split into a list of tensors of shape: (12, 1, 768).
So finally list_of_tensors contains tensors of the shape:
[torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768]),
torch.Size([12, 1, 768])]
Yes, you can directly slice it using the index and then use torch.unsqueeze() to promote the 2D tensor to 3D:
# inputs
In [6]: tensor = torch.rand(12, 512, 768)
In [7]: idx_list = [0,2,3,400,5,32,7,8,321,107,100,511]
# slice using the index and then put a singleton dimension along axis 1
In [8]: for idx in idx_list:
...: sampled_tensor = torch.unsqueeze(tensor[:, idx, :], 1)
...: print(sampled_tensor.shape)
...:
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
Alternatively, if you want it even more terse code and don't want to use torch.unsqueeze(), then use:
In [11]: for idx in idx_list:
...: sampled_tensor = tensor[:, [idx], :]
...: print(sampled_tensor.shape)
...:
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
torch.Size([12, 1, 768])
Note: there's no need to use a for loop if you wish to do this slicing only for one idx from idx_list
You can use the tensor.gather() function:
tensor = torch.rand(12, 512, 768)
ind = torch.tensor([0,2,3,400,5,32,7,8,321,107,100,511]).unsqueeze(1).unsqueeze(-1).expand(-1,-1,768) # shape (12,1,768)
tensor.gather(dim = 1, index = ind) # # shape (12,1,768)
