网络编程
位置：首页>> 网络编程>> 网络编程>> 解析Pytorch中的torch.gather()函数

解析Pytorch中的torch.gather()函数

作者：xiaoliujun1999　　发布时间：2023-01-29 23:44:40　

标签：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]])

画图说明 

解析Pytorch中的torch.gather()函数

解析Pytorch中的torch.gather()函数

官方文档


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

来源：https://blog.csdn.net/xiaoliujun1999/article/details/121292061

0
投稿

猜你喜欢

手机版 网络编程 asp之家 www.aspxhome.com