torch.gather()解释与使用

torch.gather() 是 PyTorch 中的一个函数，用于在指定维度上根据索引值从输入张量中提取数据。具体而言，torch.gather(input, dim, index) 的作用是根据 dim 维度上的 index 索引值，从 input 张量中提取对应位置的数据，并组合成一个新的张量返回。

例如，假设有一个形状为 $[3,4]$ 的输入张量 input 和一个形状为 $[3,2]$ 的索引张量 index，它们的具体数值如下：

import torch

input = torch.tensor([[1, 2, 3, 4],
                      [5, 6, 7, 8],
                      [9, 10, 11, 12]])

index = torch.tensor([[0, 2],
                      [1, 3],
                      [0, 1]])

那么，我们可以使用 torch.gather() 函数从 input 中提取对应位置的数据，如下所示：

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

输出结果如下：

tensor([[ 1,  3],
        [ 6,  8],
        [ 9, 10]])

其中，output 张量的第一行表示从 input 张量的第 0 行中提取了第 0 列和第 2 列的数据（即 [1, 3]），第二行表示从第 1 行中提取了第 1 列和第 3 列的数据（即 [6, 8]），第三行表示从第 2 行中提取了第 0 列和第 1 列的数据（即 [9, 10]）。