pytorch中torch.gather（）简单理解

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)

结果如下：