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

Question

假设我有一个sequencesshape的张量[8, 12, 2]。现在我想为每个第一维选择那个张量，这会产生一个 shape 的张量[8, 2]。维度 1 上的选择由存储在indicesshape长张量中的索引指定[8]。

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

sequences[:, indices]

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

Answer 1

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

Answer 2

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

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