图解pytorch里面的torch.gather()

在 Dim=1 的情况下应用 torch.gather()

 

 

上图显示了 torch gather() 函数在 dim=1 的二维张量上的工作。

这里索引张量的行对应于输入张量的行（用灰色阴影突出显示）。现在对于索引张量中的每个索引值，从该行和输入张量的索引中选取相应的值。

让我们一步步理解这个例子——

• 在第0行中，输入张量的第2、1、0个索引中的元素值分别为2、1和0。
• 在第 1 行中，输入张量的第 0、2 和 1 个索引中的元素值分别为 3、5 和 4。

代码实现如下所示——

在 [5] 中：

torch.gather(input=tensor2, dim=1, index = torch.tensor([[2, 1, 0], 
                                                         [0, 2, 1]]))

出[5]：

张量（[[2, 1, 0],
        [3, 5, 4]])

参考文献:

[Diagram] How to use torch.gather() Function in PyTorch with Examples - MLK - Machine Learning Knowledge