如何在pytorch中动态索引张量？

Question

例如，我得到一个张量：

tensor = torch.rand(12, 512, 768)

我得到了一个索引列表，说它是：

[0,2,3,400,5,32,7,8,321,107,100,511]

在给定索引列表的情况下，我希望从维度 2 的 512 个元素中选择 1 个元素。然后张量的大小将变为(12, 1, 768)。

有办法吗？

Answer 1

是的，您可以直接使用索引对其进行切片，然后使用torch.unsqueeze() 将2D 张量提升为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])

---

或者，如果您想要更简洁的代码并且不想使用torch.unsqueeze()，则使用：

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])

---

注意：如果您只想对来自idx_list 的一个idx 进行此切片，则无需使用for 循环

Answer 2

还有一种方法只使用 PyTorch 并使用 indexing 和 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)

当您调用tensor[:, idx_tensor, :] 时，您将得到一个形状张量：
(12, len_of_idx_list, 768)。
第二个维度取决于您的索引数量。

使用torch.split，这个张量被分割成一个张量的列表，形状为：(12, 1, 768)。

所以最后list_of_tensors 包含形状的张量：

[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])]