torch.garther函数的numpy形式

最新推荐文章于 2023-04-20 16:45:29 发布

你看见的我

最新推荐文章于 2023-04-20 16:45:29 发布

阅读量3k

点赞数 1

分类专栏：技巧学习

本文链接：https://blog.csdn.net/qq_42328228/article/details/116538936

版权

torch.gather numpy 索引张量操作数据处理

关键词由CSDN通过智能技术生成

技巧学习专栏收录该内容

2 篇文章 0 订阅

订阅专栏

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]])