pytorch中的torch.gather()函数的理解和使用

在看pytorch 时，遇到了gather函数，主要对其中dim =0，dim=1有很多困惑不清楚如何进行取数的，后面明白了点作为自己总结也希望可以帮到大家

# 选取对角线的元素
import torch as t
a = t.arange(0, 16).view(4, 4)
index = t.LongTensor([[0,1,2,3]])
b=a.gather(0, index)
print(b)

我们先了解一下这个函数以及其中的参数都代表的什么意思，

torch.gather(input, dim, index, out=None) → Tensor

input:要进行操作的tensor数据，也就是例子中的a

dim:维度

index:索引列表

函数操作就是按照给定轴dim的设置，根据index索引列表取出input数据中相对应位置的数，通过这些参数可能还是不太容易明白，其中最难理解的的就是维度dim这个参数了现在我们看几个例子，相信看完之后少许会明白一些

..............下面是dim = 0时测试........

import torch
a = torch.Tensor([[1,2,3],[4,5,6]]