以官方说明为例，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]])
官方文档
def gather(self, input, dim, index, *args, **kwargs):
For a 3-D tensor the output is specified by::
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
Args:
input (Tensor): the source tensor
dim (int): the axis along which to index
index (LongTensor): the indices of elements to gather
Example::
>>> t = torch.tensor([[1, 2], [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
tensor([[ 1, 1],
[ 4, 3]])
