pytorch中gather函数的理解。

Posted

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch中gather函数的理解。相关的知识,希望对你有一定的参考价值。

参考技术A 函数 torch.gather(input, dim, index, out=None) → Tensor
沿给定轴 dim ,将输入索引张量 index 指定位置的值进行聚合.
对一个 3 维张量,输出可以定义为:

Parameters:

使用说明举例:

Pytorch的gather用法理解

先放一张表,可以看成是二维数组

行(列)索引 索引0 索引1 索引2 索引3
索引0 0 1 2 3
索引1 4 5 6 7
索引2 8 9 10 11
索引3 12 13 14 15

看一下下面例子代码:

针对0维(输出为行形式)

>>> import torch as t
>>> a = t.arange(0,16).view(4,4)
>>> a
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
 
#选取对角线的元素
>>> index = t.LongTensor([[0,1,2,3]])
>>> a.gather(0,index)
tensor([[ 0,  5, 10, 15]])

如何理解结果呢?其实很简单,就是a.gather(0,index)中第一个0已经表明输出结果是行形式(0维),如果第一个是1说明输出结果是列形式(1维),然后按照index = tensor([[0, 1, 2, 3]])顺序作用在行上索引依次为0,1,2,3

  • a[0][0] = 0
  • a[1][1] = 5
  • a[2][2] = 10
  • a[3][3] = 15

针对0维

# 选取反对角线上的元素,注意与上面的不同
>>> index2 = t.LongTensor([[3,2,1,0]])
>>> a.gather(0,index2)
tensor([[12,  9,  6,  3]])

如何理解结果呢?同理,按照index = tensor([[3, 2, 1, 0]])顺序作用在行上索引依次为3,2,1,0:

  • a[3][0] = 12
  • a[2][1] = 9
  • a[1][2] = 6
  • a[0][3] = 3

针对1维(输出为列形式)

#选取对角线的元素

>>> index3 = t.LongTensor([[0,1,2,3]]).t()
>>> a.gather(1,index3)
tensor([[ 0],
        [ 5],
        [10],
        [15]])

如何理解结果呢?同理,按照index = tensor([[0, 1, 2, 3]])顺序作用在列上索引依次为0,1,2,3:

  • a[0][0] = 0
  • a[1][1] = 5
  • a[2][2] = 10
  • a[3][3] = 15

针对1维

选取反对角线上的元素

>>> index4 = t.LongTensor([[3,2,1,0]]).t()
>>> a.gather(1,index4)
tensor([[ 3],
        [ 6],
        [ 9],
        [12]])

如何理解结果呢?同理,按照index = tensor([[3, 2, 1, 0]])顺序作用在列上索引依次为3,2,1,0:

  • a[0][3] = 3
  • a[1][2] = 6
  • a[2][1] = 9
  • a[3][0] = 12

以上是关于pytorch中gather函数的理解。的主要内容,如果未能解决你的问题,请参考以下文章

Pytorch的gather用法理解

pytorch 笔记:gather 函数

pytorch - torch.gather 的倒数

小白学习之pytorch框架-softmax回归(torch.gather()torch.argmax())

torch.gather()之通俗易懂讲解

PyTorch中scatter和gather的用法