torch.gather的作用

torch.gather()作为PyTorch函数,它可以从一个张量中按照指定的索引来抽取元素。

其主要作用包括:

1. 从一个多维张量中按照指定的轴dim和索引索引index,抽取元素并聚集到一个新的张量中。
2. 可以实现按照指定索引对张量进行切片、重排等操作。
3. 可以和torch.scatter()配合实现在指定位置插入或复制张量元素的效果。

torch.gather()的基本语法如下:

output = torch.gather(input, dim, index)

其中:

• input: 输入张量
• dim: 沿着此轴进行gather操作
• index: 索引,指定了从input中抽取哪些元素,形状大小需要和input在dim这个轴上的形状相同
• output: 输出张量,shape为index.shape + input.shape[dim+1:]

一个简单的示例:

import torch
input = torch.tensor([[1,2],[3,4]])
index = torch.tensor([[0,0],[1,0]])
output_0 = torch.gather(input, 0, index)
print(output_0)
output_1 = torch.gather(input, 1, index)
print(output_1)

输出结果:

tensor([[1, 2],
        [3, 2]])

tensor([[1, 1],
        [4, 3]])

这里沿着输入张量的第0轴(横轴)/第1轴(列轴)进行gather,索引index指定了每一列/行抽取的元素索引,最终得到新的输出张量。

torch.gather()是一个非常有用的函数,可以实现许多gather、切片、重排类的操作。需要配合dim和index参数设定来达到不同的效果。