深度学习笔记（三）——对Tensor索引操作函数gather的解释

事先声明：本文只会对二维张量的gather操作进行介绍，三维张量的gather操作规则在csdn上的博文屡见不鲜。本文的解释是从个人的理解出发，相信解释也会对理解三维张量的操作规则起到触类旁通的作用。

gather函数的输出规则

$out[i][j]=input[index[i][j]][j],ifdim==0$
$out[i][j]=input[i][index[i][j]],ifdim==1$

第一条规则

从行的角度出发，输入的index张量按照如上规则，取出对应的输入张量的元素。
例如：一个$2\ast 3$的维度的张量，index为$[0,1,1]$，取$dim=0$，根据规则，外层循环为变量$i$，内层循环为变量$j$，且$iinrange(0,2);jinrange(0,3)$。
代入$i=0,j=1$，得到：
$out[0][1]=input[index[0][1]][1]$
$out[0][1]=input[1][1]$
即：该输出元素为输入的$2\ast 3$维度张量的第1行第1列元素。且该元素在输出张量中处在第0行第1列的位置。
如下表所示：

	0	1	2
0
1		this element

其中，$0,1,2$代表列标号，$0,1$代表行标号。

第二条规则

从列的角度出发，输入的index张量按照如上规则，取出对应的输入张量的元素。
例如：一个$2\ast 3$的维度的张量，index为$[[0,1,1],[1,1,1]]$，取$dim=1$，根据规则，外层循环为变量$i$，内层循环为变量$j$，且$iinrange(0,2);jinrange(0,3)$。
代入$i=1,j=1$，得到：
$out[1][1]=input[1][index[1][1]]$
$out[1][1]=input[1][1]$
即：该输出元素为输入的$2\ast 3$维度张量的第1行第1列元素，且该元素在输出张量中处在第1行第1列的位置。
如下表所示：

	0	1	2
0
1		this element

其中，$0,1,2$代表列标号，$0,1$代表行标号。

gather函数内部的代码机理推测

声明：下述代码仅针对原理部分编写，距离函数内部真实情况仍存在较大差距，且下述代码的严谨性不够，故仅供理解gather的核心规则。

def gather(input, dim, index):
# 这里的dim要求取0或1
	out = []
	m = input.size()[0] # size函数是torch的方法
	n = input.size()[1]
	for i in range(m):
		for j in range(n):
			if dim == 0:
				out [i] [j] = input [index [i] [j] ] [j]
			if dim == 1:
				out [i] [j] = input [i] [index [i] [j] ]
	return out

代码示例

与上一篇博文内容相同，这里再次展示一遍。

import torch
# 设置一个随机种子
torch.manual_seed(100)
# 生成一个形状为2*3的矩阵
x = torch.randn(2, 3)
print(x)
# 获取指定索引对应的值
index = torch.LongTensor([[0, 1, 1]])
print(torch.gather(x, 0, index))
index = torch.LongTensor([[0, 1, 1], [1, 1, 1]])
a = torch.gather(x, 1, index)
print(a)

输出结果

参考文献

吴茂贵,郁明敏,杨本法,李涛,张粤磊. Python深度学习（基于Pytorch）. 北京:机械工业出版社,2019.