pytorch中torch.gather()简单理解

博客介绍了在PyTorch中从输入张量按指定维度进行索引采集操作,返回新张量。同时指出该函数主要问题在dim维度上,如dim=0沿行采集,dim=1沿列采集,并通过具体例子说明了不同dim值下的采集规则。

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

1.作用

从输入张量中按照指定维度进行索引采集操作,返回值是一个新的张量,形状与 index 张量相同,根据指定的索引从输入张量中采集对应的元素。

2.问题

该函数的主要问题主要在dim维度上,dim=0 表示沿着第一个维度(行)进行索引采集,而 dim=1 表示沿着第二个维度(列)进行索引采集。
简单讲:dim=0,将在行上进行采集,行数不变,在列上取值,如下图中的例子torch.gather(input, dim=0, index=index),当dim=0时,[0,1]中0对应第一行第一列也就是1,1对应着第2行第2列(1在index的坐标为(1,2),dim=0,所以不用看index的行坐标的,只管纵坐标,也就是第2列。而此时值为1代表值input的行【需要+1,下标是从0开始的】,也就是第2行,值为4)也就是4;[1,0]中的1代表第2行第1列也就是3,0也就是第2列的第一行数据(此时index的0的坐标为(2,2),因为dim=0,也就不用看横坐标,也就是第二列。所以此时的0代表源input的行坐标,也就是第一行)也就是2

input = torch.tensor([[1, 2], [3, 4], [5, 6]])
index = torch.tensor([[0, 1], [1, 0]])

result = torch.gather(input, dim=0, index=index)
result_colum = torch.gather(input, dim=1, index=index)
print("result:",result)
print("result_colum:",result_colum)

结果如下:
在这里插入图片描述

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值