理解 PyTorch 中的 gather 函数

我会带着你远行 2022-09-09 06:03 155阅读 0赞

好久没更新博客了，最近一直在忙，既有生活上的也有工作上的。道阻且长啊。

今天来水一文，说一说最近工作上遇到的一个函数：[`torch.gather()`][torch.gather] 。

## 文字理解 ##

我遇到的代码是 NLP 相关的，代码中用 `torch.gather()` 来将一个 tensor 的 shape 从 `(batch_size, seq_length, hidden_size)` 转为 `(batch_size, labels_length, hidden_size)` ，其中 `seq_length >= labels_length` 。

`torch.gather()` 的官方解释是

> Gathers values along an axis specified by *dim*.

就是在指定维度上 gather value。那么怎么 gather、gather 哪些 value 呢？这就要看其参数了。

`torch.gather()` 的必填也是最常用的参数有三个，下面引用官方解释：

*  `input` (*[Tensor][]*) – the source tensor
 *  `dim` (*[int][]*) – the axis along which to index
 *  `index` (*LongTensor*) – the indices of elements to gather

所以一句话概括 gather 操作就是：**根据 `index` ，在 `input` 的 `dim` 维度上收集 value**。

具体来说，`input` 就是源 tensor，等会我们要在这个 tensor 上执行 gather 操作。如果 `input` 是一个一维数组，即 flat 列表，那么我们就可以直接根据 `index` 在 `input` 上取了，就像正常的列表/数组索引一样。但是由于 `input` 可能含有多个维度，是 N 维数组，所以我们需要知道在哪个维度上进行 gather，这就是 `dim` 的作用。

对于 `dim` 参数，一种更为具体的理解方式是**替换法**。假设 `input` 和 `index` 均为三维数组，那么输出 tensor 每个位置的索引是列表 `[i, j, k]` ，正常来说我们直接取 `input[i, j, k]` 作为 输出 tensor 对应位置的值即可，但是由于 `dim` 的存在以及 `input.shape` 可能不等于 `index.shape` ，所以直接取值可能就会报 `IndexError` 。所以我们是将索引列表的相应位置替换为 `dim` ，再去 `input` 取值。如果 `dim=0` ，我们就替换索引列表第 0 个值，即 `[dim, j, k]` ，依此类推。Pytorch 的官方文档的写法其实也是这个意思，但是看这么多个方括号可能会有点懵：

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

但是可能你还有点迷糊，没关系接着看下面的直观理解部分，然后再回来看这段话，结合着看，相信你很快能明白。

由于我们是按照 `index` 来取值的，所以最终得到的 tensor 的 shape 也是和 `index` 一样的，就像我们在列表上按索引取值，得到的输出列表长度和索引相等一样。

## 直观理解 ##

为便于理解，我们以一个具体例子来说明。我们使用反推法，根据 `input` 和输出推参数。这应该也是我们平常自己写代码的时候遇到比较多的情况。

假设 `input` 和我们想要的输出 `output` 如下：

>>> input_tensor = torch.arange(24).reshape(2, 3, 4)
    tensor([[[ 0,  1,  2,  3],
             [ 4,  5,  6,  7],
             [ 8,  9, 10, 11]],
    
            [[12, 13, 14, 15],
             [16, 17, 18, 19],
             [20, 21, 22, 23]]])
    >>> output_tensor  # shape: (2, 2, 4)
    tensor([[[ 0,  1,  2,  3],
             [ 8,  9, 10, 11]],
    
            [[12, 13, 14, 15],
             [20, 21, 22, 23]]])

即，我们想让 shape 为 `(2, 3, 4)` 的 `input_tensor` 变成 shape 为 `(2, 2, 4)` 的 `output_tensor` ，丢弃维度 1 的第 2 个元素，即 `[ 4, 5, 6, 7]` 和 `[16, 17, 18, 19]` 。

我们应用替换法，重点是找出来 `dim` 和 `index` 的值。始终记住 `index` 和 `output_tensor` 的 shape 是一样的。

从 `output_tensor` 的第一个位置开始，由于 `output_tensor[0, 0, :] = input_tensor[0, 0, :]` ，所以此时 `[i, j, k]` 是一样的，我们看不出来 `dim` 应该是多少。

下一行 `output_tensor[0, 1, 0] = input_tensor[0, 2, 0]` ，这里我们看到维度 1 发生了变化，1 变成了 2，所以 `dim` 应该是 1，而 `index` 应为 2， `index_tensor[0, 1, 0]=2` 。

此时 `dim` 已经明确。同理，`output_tensor[0, 1, 1] = input_tensor[0, 2, 1]` ，`index_tensor[0, 1, 1]=2` ，依此类推，得到 `index_tensor[0, 1, :] = 2` 。同时也可以明确 `index_tensor[0, 0, :] = 0` 。

所以

>>> dim = 0
    >>> index_tensor
    tensor([[[0, 0, 0, 0],
             [2, 2, 2, 2]],
    
            [[0, 0, 0, 0],
             [2, 2, 2, 2]]])

简单可描述如下图：  
![torch.gather() 执行过程][torch.gather_]  
为描述方便，假如我们把输入看作是 6 行，从上到下依次是 0-5。那么从事后诸葛亮的角度讲，输出相当于是把第 1 和第 4 行“抽掉”。如果输出和输入一样，那么原本的 `index_tensor` 就是如下：

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

“抽掉”后， `index_tensor` 也相应“抽掉”，那么就得到我们想要的结果了。而且由于这个“抽掉”的操作是在维度 1 上进行的，那么 `dim` 自然是 1。

`numpy.take()` 和 `tf.gather` 貌似也是同样功能，就不细说了。

## Reference ##

*  [torch.gather — PyTorch 1.9.0 documentation][torch.gather]
 *  [numpy.take — NumPy v1.21 Manual][numpy.take _ NumPy v1.21 Manual]
 *  [tf.gather | TensorFlow Core v2.6.0][tf.gather _ TensorFlow Core v2.6.0]

## END ##

[torch.gather]: https://pytorch.org/docs/stable/generated/torch.gather.html
[Tensor]: https://pytorch.org/docs/stable/tensors.html#torch.Tensor
[int]: https://docs.python.org/3/library/functions.html#int
[torch.gather_]: /images/20220829/f6ad04883e3e4d75b111685fa7425c9f.png
[numpy.take _ NumPy v1.21 Manual]: https://numpy.org/doc/stable/reference/generated/numpy.take.html
[tf.gather _ TensorFlow Core v2.6.0]: https://www.tensorflow.org/api_docs/python/tf/gather