Pytorch中的torch.gather()函数怎么用

这篇文章将为大家详细讲解有关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]])

画图说明 

官方文档

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]])

关于“Pytorch中的torch.gather()函数怎么用”这篇文章就分享到这里了，希望以上内容可以对大家有一定的帮助，使各位可以学到更多知识，如果觉得文章不错，请把它分享出去让更多的人看到。