参数说明
以官方说明为例,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]])