pytorch scatter和scatter_详解

0. Introduction

scatter() 和 scatter_() 的作用是一样的，只不过 scatter() 不会直接修改原来的 Tensor，而 scatter_() 会

PyTorch 中，一般函数加下划线代表直接在原来的 Tensor 上修改

1. 定义

scatter(dim, index, src) 的参数有 3 个

dim：沿着哪个维度进行索引
index：用来 scatter 的元素索引
src：用来 scatter 的源元素，可以是一个标量或一个张量

这个 scatter 可以理解成放置元素或者修改元素

简单说就是通过一个张量 src 来修改另一个张量，哪个元素需要修改、用 src 中的哪个元素来修改由 dim 和 index 决定

官方文档给出了 3维张量 的具体操作说明，如下所示

self [index[i][j][k]] [j] [k] = src[i][j][k]  # if dim == 0
self[i] [index[i][j][k]] [k] = src[i][j][k]  # if dim == 1
self[i] [j] [index[i][j][k]] = src[i][j][k]  # if dim == 2

如果是二维的例子，则应该对应下面的情况：

y [ index[i][j] ] [j] = src[i][j] #if dim==0
y[i] [ index[i][j] ]  = src[i][j] #if dim==1 

2. 详解

例1

源tensor的每个元素，都按照 index 被scatter（可以理解为填充）到目标tensor中。

1. 考虑最简单的目标：把x直接塞入目标tensor的前两行，dim = 0时，index应该是

   [[0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1]]
   

2. 把x直接塞入目标tensor的后两行，dim = 0，index应该是

   [[1, 1, 1, 1, 1],
    [2, 2, 2, 2, 2]]
   

所以，index为几，就把对应位置的元素放入目标tensor的第几行（dim=0时，列不变），大概就是这么个意思。以此类推，dim=1时，“直接塞入”对应的index就是

[[0,1,2,3,4],
 [0,1,2,3,4]]

例2

更为复杂的例子：

import torch

a = torch.arange(1,11).reshape(2,5).float()
# tensor([[1, 2, 3, 4, 5],
#          6, 7, 8, 9, 10]])

index=torch.LongTensor([[1, 2, 1, 1, 2], [2, 0, 2, 1, 0]])
# tenso[[1, 2, 1, 1, 2], 
#		[2, 0, 2, 1, 0]]

b = torch.zeros(3, 5)
b_= b.scatter(dim=0, index=index, src=a)

# tensor([[0, 7, 0, 0, 10],
#         [1, 0, 3, 9, 0],
#         [6, 2, 8, 0, 5]])

过程解析：

1. 把a中的元素一个个放入到b中
2. 先取a[0][0]=1, 因为参数dim=0，所以在b中列索引不变，我们要在b选择1这个元素所对应的行。index第一个元素是1，对应b中第二行，所以b[1][0]=1
3. 再取a[0][1]=2，列索引不变，index对应的第二个元素是2，所以行索引是2。故b[2][1]=2
4. 重复上述操作直至完成全部元素的放置

技巧：
在寻找元素和index之间的匹配关系时，可以直接按照对应位置来。最终从a选取到b中的元素由index元素的个数决定，即index有多少个元素，a就有多少个元素被选取。

Reference：

torch.scatter_直观理解官网示例
PyTorch笔记之 scatter() 函数