Pytorch中的torch.gather()函数

参数说明

以官方说明为例，gather()函数需要三个参数，输入input，维度dim，以及索引index

input必须为Tensor类型

dim为int类型，代表从哪个维度进行索引

index为LongTensor类型

举例说明

input=torch.tensor([[1,2,3],[4,5,6]]) #作为输入

index1=torch.tensor([[0,1,1],[0,1,1]]) #作为索引矩阵

# dim=0时,按列进行索引
print (torch.gather(input,dim=0,index=index1))

# dim=1时,按行进行索引
print (torch.gather(input,dim=1,index=index1))

 结果如下图所示：

# 按列进行索引
tensor([[1, 5, 6],
        [4, 2, 6]])

# 按行进行索引
tensor([[1, 2, 2],
        [5, 4, 5]])

画图说明