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
PyTorch踩坑:gather函数

上一篇:这个meta标签会让华为手机自带浏览器无法粘贴手机收到的验证码信息


下一篇:KAL1 LINUX 官方文档之kali nethunter文档 --- 在Gemini PDA上安装NetHunter