Pytorch 中 gather 函数讲解

官方解读分析

该函数的功能为：沿着 dim 指定的轴收集值

torch.gather(input, dim, index, out=None) → Tensor

    Gathers values along an axis specified by dim.

    For a 3-D tensor the output is specified by:

    out[i][j][k] = input[index[i][j][k]][j][k]  # dim=0
    out[i][j][k] = input[i][index[i][j][k]][k]  # dim=1
    out[i][j][k] = input[i][j][index[i][j][k]]  # dim=2

    Parameters:	

        input (Tensor) – The source tensor
        dim (int) – The axis along which to index
        index (LongTensor) – The indices of elements to gather
        out (Tensor, optional) – Destination tensor

    Example:

    >>> t = torch.Tensor([[1,2],[3,4]])
    >>> torch.gather(t, 1, torch.LongTensor([[0,0],[1,0]]))
     1  1
     4  3
    [torch.FloatTensor of size 2x2]

我们输出的是一个 tensor 数组，因为函数的作用是收集值，且收集的值是遍历index 得到的。因此输出数组的大小与 index 大小相同

总结来说，就是out[][][] = input[][][]。out 的参数就是位置的索引，input 的参数，第dim 维由 index 中的值确定，其余参数与 out 相同

• dim = 0时，out[i][j][k] = input[ [index[i][j][k]] [j][k]]。这代表什么意思呢？

  当我们遍历 index 来收集输出结果的时候，当前位置的输出 out[当前位置]，是 input 中的一个数。这个数的第零维是 index 中当前位置对应的数，其余维度的参数是当前位置的索引

• dim = 其他值时同理

以上面的例子为说明：

dim = 1，代表，第 1 维的参数由 index 中的数控制，第 0 维参数，out 与 input 相同

• out[0][0] = t[0][index[0][0]] = t[0][0] = 1
• out[0][1] = t[0][index[0][1]] = t[0][0] = 1
• out[1][0] = t[1][index[1][0]] = t[1][1] = 4
• out[1][1] = t[1][index[1][1]] = t[1][0] = 3

小例子

import torch
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
index = torch.tensor([[1, 0, 1],[0, 0, 0]])
print(a)
print(torch.gather(a, 1, index))
print(torch.gather(a, 0, index))

tensor([[1, 2, 3],
        [4, 5, 6]])
        
tensor([[2, 1, 2],
        [4, 4, 4]])
        
tensor([[4, 2, 6],
        [1, 2, 3]])

我们先分析 torch.gather(a, 1, index)。因为dim = 1，因此第 1 维参数由 index 中的数确定。这与官方例子很类似，不用过多讲解

那么 dim = 0 的呢？这时，第 0 维由 index 中的数确定

• out[0][0] = a[index[0][0]][0] = a[1][0] = 4
• out[0][1] = a[index[0][1]][1] = a[0][1] = 2
• out[0][2] = a[index[0][2]][2] = a[1][2] = 6

剩下的不再推导，由上例容易得出。

收集的数就是：遍历 index ，遍历到的数作为一个参数，剩下的参数就是当前数的位置。遍历到的数作为哪个参数，由 dim 决定