使用torch.gather()的方式有两种,第一种为直接使用torch.gather().第二种为SourceTensor.gather().容易出现疑惑的地方为dim的取值.dim的取值范围为{-2,0,1}.下面分情况进行讨论(在二维平面上讨论).
通过上图中的代码执行结果可以看出,当dim为0时,y中的数据值表示行索引,而y中的数据索引表示列索引.具体到上图中的情况即是,y的第n列数据yn,表示索引值[yn][n-1],例如y中的第3列数据0,表示索引值[0][2],gather()的输出,即是y_hat[0][0],y_hat[1][1],y_hat[0][2].值得注意的是,在调用gather()时应注意维度匹配和数组越界问题,例如,在此处y_hat的维度为2×3,即两行三列,由于y中的数据索引表示列索引,因此y的列数也必须为3,而y中的数据值表示行索引,因此y的数据值不可以大于或者等于2.
观察上图中的代码执行结果可以看出,dim=1时,索引情况刚好与为0时相反,y中的数据索引表示行索引,而y中的数据值表示列所引.具体到上图中的情况即是,y的第n行数据yn,表示索引值[n-1][yn],例如y中的第二行数据2,表示索引值[1][2].维度匹配以及取值范围的判断方法与dim=0时一致.
在二维平面上,dim=-2时,index的取值以及最终输出总于dim=0时一致,具体原因目前尚未弄清楚.
总的来说,调用source.gather(dim,index)函数时,当dim=0时,index的列编号将会直接用到输出结果中;当dim=1时,index的行编号将会直接用到输出结果中.
