torch.garther函数
torch.garther函数的具体原理可参见这里。
以二维数组为例,
- 下述代码是行索引的例子;
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]])
- 下述代码是列索引的例子;
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形式
- 下述代码是行索引的例子;
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]])
- 下述代码是列索引的例子;
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]])