torch.gather()用法

torch.gather()用法

在torch文档官网看了很久关于torch.gather()看了很久，终于搞懂其用法，但是经常会忘记，因此今天总结此用法；

先看官网解释：

从官网解释来看，主要3个主要参数，分别是，input， dim, index。字面意思，input数据（可以是tensor矩阵，可以是tensor向量），dim表示维度，index为索引参数；

接下来看看官网例子：
例子很简单哈，先来解释一下官网例子：

相信这个图示加上官网给的解释，应该大家明白了哈。
可能有人会询问这个用途是什么？？

那当然也很简单，比如给你一个2维矩阵H*W，你想知道dim=1维度某个位置的元素值，就用这个很方便了。换句话说，如果你想知道第一行，第1个元素，第二行第0个元素，那只需要把index=[1,0]， dim = 1就好了。

output = torch.gather(input, dim = 1, index = torch.tensor([1, 0]))