PyTorch踩坑：gather函数

整理自官方文档：https://www.oschina.net/action/GoToLink?url=https%3A%2F%2Fpytorch.org%2Fdocs%2Fstable%2Fgenerated%2Ftorch.gather.html%3Fhighlight%3Dgather%23torch.gather

0x01 背景

在做强化学习DDQN实（CtrlC）验（CtrlV）的时候遇到了一些不是很理解的函数，这里做一些解读。

0x02 gather函数

gather，直译为聚合，聚集。

先跑两个例子示范一下它会出来什么东西：

首先建立两个tensor：a和b

a = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
tensor([[1, 2, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 1, 0]])

下面展示一下效果：

>>> a
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15]])
>>> b
tensor([[1, 2, 0, 0, 0],
        [0, 1, 0, 0, 0],
        [0, 0, 0, 1, 0]])
>>> a.gather(0, b)
tensor([[ 6, 12,  3,  4,  5],
        [ 1,  7,  3,  4,  5],
        [ 1,  2,  3,  9,  5]])
>>> a.gather(1, b)
tensor([[ 2,  3,  1,  1,  1],
        [ 6,  7,  6,  6,  6],
        [11, 11, 11, 12, 11]])

gather函数其实是就相当于是对原来的tensor做一个选取和替换。

第一个参数是dim，表明我们要在哪个维度上做选择（比如是在矩阵的行还是列上做选择）。

第二个参数是索引，不一定非得和原来的a的形状保持一致。

下面，我们直接用一张网图来说明gather函数的工作原理。

（本来打了很多字，但是还是说不清楚QAQ