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)
结果如下: