torch.gather函数的理解

建议先阅读官方文档,拿笔跟着给出的公式推导一次。

torch.gather官方文档

gather函数的定义为:
torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
不常用的暂时不关注,于是函数常使用的样子如下:
torch.gather(input, dim, index)
函数的大致功能,给出input,根据dim和index确认从input中取出的数据内容,和最终输出的shape
1.确定输出的shape,输出的shape跟index的shape一致
例如index是一个2x3的tensor,那么输出就是一个2x3的tensor
2.确定输出的tensor的内容:根据input,dim,index三者确定
dim确定的是在input中取数据的时候,那个维度是从index中取数据来确定索引,而不是直接顺序索引:
举例
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2

out[i][j] = input[index[i][j][k]][j] # if dim == 0
out[i][j] = input[i][index[i][j][k]] # if dim == 1

其实这个函数一般是结合其他的函数使用的,比如torch.max,torch.min。比如调用max函数如下:
torch.max(a,dim=dim_num)
那么此时返回的是一个tuple,包含了value和index_list
那么有value = torch.gather(a, dim=dim_num, index=index_list)

©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容