tensor.gather()的使用(pytorch)

页页读

已于 2022-05-25 09:23:11 修改

阅读量8.2k

点赞数 7

分类专栏： pytorch 文章标签： pytorch torch.gather tensor.gather

于 2020-10-23 15:41:06 首次发布

本文链接：https://blog.csdn.net/u014386899/article/details/109243383

版权

pytorch 专栏收录该内容

18 篇文章

订阅专栏

本文详细介绍了PyTorch中的tensor.gather()函数，该函数用于在指定维度上根据提供的索引选取元素。通过举例说明了如何使用此函数，返回与索引相同形状的tensor。示例中展示了如何在二维tensor上进行操作，选取特定位置的元素。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tensor.gather()的使用(pytorch)

tensor.gather(dim, indexs)

功能：在dim维度上，按照indexs所给的坐标选择元素，返回一个和indexs维度相同大小的tensor。
它和torch.gather功能是一样的。
在这里插入图片描述
torch.gather()官方文档

注意：这里indexs必须也是Tensor，并且维度数与input相同（len(input.shape)=len(indexs.shape)）

例子：

>>>import torch
>>>a = torch.Tensor([[1,2,3,4,5,6],[0,1,2,3,4,5]])
>>>a
Out[4]: 
tensor([[1., 2., 3., 4., 5., 6.],
        [0., 1., 2., 3., 4., 5.]])
>>>
>>>b = a.gather(1, torch.tensor([[5,4,3,2,1,0],[0,1,2,3,4,5]]))
>>>b
Out[9]: 
tensor([[6., 5., 4., 3., 2., 1.],
        [0., 1., 2., 3., 4., 5.]])