pytorch的gather

1. 官网解释

$out[i][j][k]=input[[index[i][j][k]][j][k],ifdim==0$
以此类推

意思就是说，假设我要求$out[i][j]$位置的值，如果dim=0，那么$i$就会被$index[i][j]$代替。而如果dim=1，那么$j$就会被$index[i][j]$代替。其他类推。

2. 一个例子

假如我有一个这样的tensor

t = torch.tensor([[1,2],[3,4]])
out = torch.gather(t,0,index=torch.tensor([[0,0],[1,0]]))
print(out)

此时dim = 0，所以
$out[0][0]=input[index[0][0]][0]=input[0][0]=1$
$out[0][1]=input[index[0][1]][1]=input[0][1]=2$
$out[1][0]=input[index[1][0]][0]=input[1][0]=3$
$out[1][1]=input[index[1][1]][1]=input[0][1]=2$
所以结果应该是$[[1,2],[3,2]]$

3. 代码验证

4. 我悟了

其实可以这么看，把值这么一摆：