torch.garther函数的numpy形式

torch.garther函数

torch.garther函数的具体原理可参见这里。
以二维数组为例，

1. 下述代码是行索引的例子；

tensor_0 = tensor([[ 3,  4,  5],
                   [ 6,  7,  8],
                   [ 9, 10, 11]])，
index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(axis=0, index)
output: tensor([[9, 7, 5]])

2. 下述代码是列索引的例子；

tensor_0 = tensor([[ 3,  4,  5],
                   [ 6,  7,  8],
                   [ 9, 10, 11]])，
index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(axis=1, index)
output: tensor([[5, 7, 9]])

经过探索后发现，此种索引方式似乎只支持torch.tensor数据，而在numpy数组中，也需要类似的索引方式。然而，使用list列表迭代方式寻找则过于繁琐，这里根据torch.gather()内置函数的原理，设计适合numpy数组的索引方法。

torch.garther函数的numpy形式

1. 下述代码是行索引的例子；

tensor_0 = np.array([[ 3,  4,  5],
                   [ 6,  7,  8],
                   [ 9, 10, 11]])，
index = np.array([[2, 1, 0]]).squeeze()
arg = np.arange(0, tensor_0.shape[1])   # 行采用index索引，列采用自然升序索引
tensor_1 = tensor_0[index, arg]
output: tensor([[9, 7, 5]])

2. 下述代码是列索引的例子；

tensor_0 = tensor([[ 3,  4,  5],
                   [ 6,  7,  8],
                   [ 9, 10, 11]])，
index = torch.tensor([[2, 1, 0]]).squeeze()
arg = np.arange(0, tensor_0.shape[0])   # 行采用自然升序索引，列采用index索引
tensor_1 = tensor_0[arg, index]
output: tensor([[5, 7, 9]])