pytorch学习（二）：tensor的索引操作--- gather和scatter_

最新推荐文章于 2025-07-17 15:28:35 发布

原创最新推荐文章于 2025-07-17 15:28:35 发布 · 2.6k 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch #人工智能

pytorch学习专栏收录该内容

3 篇文章

订阅专栏

本文深入解析了PyTorch中gather与scatter_函数的使用方法及应用场景，gather函数根据指定索引在特定维度上选取数据，scatter_则为数据放回操作，通过实例展示了如何在二维张量中进行对角线、反对角线元素的选取与放置，适合深度学习开发者学习。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1.gather函数

gather(input, dim, index)：根据index，在dim维度上选取数据，输出的size与index一样。
gather是一个比较复杂的操作，对一个二维tensor，输出的每个元素如下：

out[i][j] = input[index[i][j]][j]   #dim = 0
out[i][j] = input[i][index[i][j]]   #dim = 1

例：

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


In:  #选取反对角线上的元素
     index = t.LongTensor([[3, 2, 1, 0]]).t()
     a.gather(1, index)
Out: tensor([[ 3],
        [ 6],
        [ 9],
        [12]])


In:  #选取两个对角线上的元素
     index = t.LongTensor([[0, 1, 2, 3], [3, 2, 1, 0]]).t()
     b = a.gather(1, index)
     b
Out: tensor([[ 0,  3],
        [ 5,  6],
        [10,  9],
        [15, 12]])

2.scatter_函数

与gather相对应的逆操作是scatter_，gather把数据从input中按index取出，而scatter_是把取出的数据再放回去。注意scatter_函数是inplace操作。
例：

In:  #把两个对角线元素放回到指定位置
    c = t.arange(2, 18).view(4, 4)
    c.scatter_(1, index, b)
Out: tensor([[ 0,  3,  4,  3],
        [ 6,  5,  6,  9],
        [10,  9, 10, 13],
        [12, 15, 16, 15]])