整理自官方文档: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