如何在 pytorch 的维度上选择单个索引？答案

【问题标题】：How can I select single indices over a dimension in pytorch?如何在 pytorch 的维度上选择单个索引？
【发布时间】：2019-05-27 22:02:23
【问题描述】：

假设我有一个形状为 [8, 12, 2] 的张量 sequences。现在我想为每个第一维选择该张量，从而产生形状为[8, 2] 的张量。维度 1 上的选择由存储在形状为 [8] 的长张量 indices 中的索引指定。

我试过这个，但是它为sequences 中的每个第一个维度选择indices 中的每个索引，而不是只选择一个。

sequences[:, indices]

如何在没有缓慢而丑陋的for 循环的情况下进行此查询？

【问题讨论】：

标签： pytorch tensor

【解决方案1】：

sequences[torch.arange(sequences.size(0)), indices]

【讨论】：

谢谢，正是我想要的。最后结果很容易:)。

【解决方案2】：

你好像在找torch.gather:

torh.gather(sequences, dim=1, index=indices)

【讨论】：

在我指定的情况下似乎不起作用。我需要如何重塑我的指数？
试试这个torh.gather(sequences, dim=1, index=indices.unsqueeze(1))
那根本行不通。 sequences = torch.randn(8, 12, 2)indices = torch.zeros(12, dtype=torch.long)torch.gather(sequences, dim=1, index=indices.unsqueeze(1))> RuntimeError: invalid argument 4: Index tensor must have same dimensions as input tensor at /Users/distiller/project/conda/conda-bld/pytorch_1565272526878/work/aten/src/TH/generic/THTensorEvenMoreMath.cpp:453

【解决方案3】：

torch.index_select 比 torch.gather 更容易解决您的问题，因为您不必调整索引的尺寸。 Indeces 必须是张量。适合你的情况

indeces = [0,2]
a = torch.rand(size=(3,3,3))
torch.index_select(a,dim=1,index=torch.tensor(indeces,dtype=torch.long))

【讨论】：

【解决方案4】：

torch.gather 应该可以做到这一点，但是您需要先将索引张量转换为

unsqueeze 匹配输入张量的维数
repeat_interleave 匹配最后一维的大小

以下是根据您的描述的示例：

# original indices dimension [8]
# after first unsueeze, dimension is [8, 1]
indices = torch.unsqueeze(indices, 1)

# after second unsueeze, dimension is [8, 1, 1]
indices = torch.unsqueeze(indices, 2)

# after repeat, dimension is [8, 1, 2]
indices = torch.repeat_interleave(indices, 2, dim=2)

# now you have the right dimension for torch.gather
# don't forget to squeeze the redundant dimension
# result has dimension [8, 2] 
result = torch.gather(sequences, 1, indices).squeeze()

【讨论】：

【解决方案5】：

可以使用torch.Tensor.gather来完成

sequences = torch.randn(8,12,2)
# defining the indices as a 1D Tensor of random integers and reshaping it to use with Tensor.gather
indices = torch.randint(sequences.size(1),(sequences.size(0),)).unsqueeze(-1).unsqueeze(-1).repeat(1,1,sequences.size(-1))
# indices shape: (8, 1, 2)
output = sequences.gather(1,indices).squeeze(1)
# output shape: (8, 2)

【讨论】：