通过索引列表选择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）指定为第二个参数即可。这是为一维输入张量指定的唯一有效维度。

Answer 2

index_select

是一个合理的选择，但它将返回所选条目的副本。如果您只想对原始张量进行一些切片，您可以使用它来访问目标项目：

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

如果你有一个二维张量：

t = torch.rand(7,2)
output = t[(0,3,4),:]

三维张量：

t = torch.rand(7,2,2)
output = t[(0,3,4),:,:]