python - 通过索引列表选择pytorch张量元素

Question

我想我有一个非常简单的问题。让我们采用以下长度为 6

的张量

t = torch.tensor([10., 20., 30., 40., 50., 60.])

现在我只想访问特定索引处的元素，例如 [0, 3, 4] 处的元素。所以我想回来

# exptected output 
tensor([10., 40., 50.])

我找到了torch.index_select这对于二维张量非常有效，例如维度 (2, 4)，但不适用于给定的 t。

如何在不使用 for 循环的情况下基于一维张量中给定的索引列表访问一组元素？

Answer 1

事实上，您可以使用 index_select 来实现此目的:

t = torch.tensor([10., 20., 30., 40., 50., 60.])
output = torch.index_select(t, 0, torch.LongTensor([0, 3, 4]))
# output: tensor([10., 40., 50.])

您只需将尺寸 (0) 指定为第二个参数即可。这是为一维输入张量指定的唯一有效维度。