Numpy/Torch：使用维度上的索引选择元素

Question

给定如下所示的数组：

np.arange(12).reshape(4,3)
Out[119]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

我想使用索引列表[0, 2, 1, 2]从每一行中选择一个元素，以创建[0, 5, 7, 11]的4x1数组

有没有简单的方法来建立索引。我能看到的最接近的方法是pytorch中的gather方法

Answer 1

>>> import torch
>>> import numpy as np
>>> s = np.arange(12).reshape(4,3)
>>> s = torch.tensor(s)
>>> s
tensor([[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
>>> idx = torch.tensor([0, 2, 1, 2])
>>> torch.gather(s,-1 ,idx.unsqueeze(-1))
tensor([[ 0],
        [ 5],
        [ 7],
        [11]])

torch.gather(s,-1 ,idx.unsqueeze(-1))