如何在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

还有一种方法只使用 PyTorch 并使用索引和 避免循环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])]