torch.scatter与torch_scatter库使用整理

最近在做图结构相关的算法，scatter能把邻接矩阵里的信息修改，或者把邻居分组算个sum或者reduce，挺方便的，简单整理一下。

torch.scatter 与 tensor._scatter

Pytorch自带的函数，用来将作为src的tensor根据index的描述填充到input中，形式如下：

ouput = torch.scatter(input, dim, index, src)
# 或者是
input.scatter_(dim, index, src)

两个方法的功能是相同的，而带下划线的_scatter方法是将原tensor input直接修改了，不带的则会返回一个新的tensor output，input不变。

其中dim决定index对应值是沿着哪个维度进行修改。而src为数据来源，当其为tensor张量时，shape要和index相同，这样index中每个元素都能对应src中对应位置的信息。

理解scatter方法主要是要理解index实现的src和input之间的位置对应关系，举个例子：

dim = 0
index = torch.tensor(
	[[0, 2, 2], 
	[2, 1, 0]]
)

dim为0时，遵循的映射原则为：input[index[i][j]][j] = src[i][j].
也就是说，将位置 (i, j) 中dim对应的位置改为 index[i][j] 的值。如位置(1,0)，index[1][0]为2，则映射后的位置为(2,0)，意味着input中(2,0)的位置被更改为src中(1,0)位置的值。

我个人形象理解是这些值会沿着dim方向滑动，上面例子中src[1][0]位置的值滑到2，成为input中的新值，这样理解起来更形象一点。

基本理解了上面这个例子，多维情况和不同dim的情况都可以类推了。

需要注意：src和input的dtype需要相同，不然会报Expected self.dtype to be equal to src.dtype，不一样就先转换再使用。

t = torch.arange(6).view(2, 3)
t = t.to(torch.float32)
print(t)

output = torch.scatter(torch.zeros((3, 3)), 0, torch.tensor([[0, 2, 2], [2, 1, 0]]), t)
print(torch.zeros((3, 3)).scatter_(0, torch.tensor([[0, 2, 2], [2, 1, 0]]), t))

输出：

tensor([[0., 1., 2.],
        [3., 4., 5.]])
tensor([[0., 0., 5.],
        [0., 4., 0.],
        [3., 1., 2.]])

torch_scatter库

这个第三方库对矩阵的分组处理这个概念做了更进一步的封装，通过index来指定分组信息，将元素分组后进行对应处理，最基础的scatter方法形式如下：

torch_scatter.scatter(src, index, dim, out, dim_size, reduce)

src: 数据源
index：分组序列
dim：分组遵循的维度
out：输出的tensor，可以不指定直接让函数输出
dim_size：out不指定的时候，将输出shape变为该值大小；dim_size也不指定，就根据计算结果来
reduce：分组的操作，包括sum，mul，mean，min和max操作

这个方法理解关键在index的分组方法，举个例子：

dim = 1
index = torch.tensor([[0, 1, 1]])

torch_scatter.scatter对index的顺序是没有特定规定的，相同数字对应的元素即为一组。比如例子中，维度1上的第0个元素为一组，第1和2元素为另一组。这样，按照分组进行reduce定义的计算即可获得输出。如：

t = torch.arange(12).view(4, 3)
print(t)
t_s = torch_scatter.scatter(t, torch.tensor([[0, 1, 1]]), dim=1, reduce='sum')
print(t_s)

输出：

tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
tensor([[ 0,  3],
        [ 3,  9],
        [ 6, 15]])

可以看出，每行的后两个元素求了和，与index定义相同。

要注意的是，index的shape[0]为1时，会自动对dim对应的维度上每一层进行相同的分组处理，如上例所示，index大小为(1, 3)，即对src的三行数据都进行了分组处理。

而另一种分组方式，如需要每行分组不同，则需要index的shape和src的shape相同，如下例：

t = torch.arange(12).view(4, 3)
print(t)
t_s = torch_scatter.scatter(t, torch.tensor([[0, 1, 1], [1, 1, 0], [0, 1, 1], [1, 1, 0]]), dim=1, reduce='sum')
print(t_s)

输出：

tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
tensor([[ 0,  3],
        [ 5,  7],
        [ 6, 15]])

shape不相同时，则会报错提示：RuntimeError: The expanded size of the tensor (3) must match the existing size (2) at non-singleton dimension 0.

同时，该库还给出了另外两种方法，分别为torch_scatter.segment_coo和torch_scatter.segment_csr.

torch_scatter.segment_coo

torch_scatter.segment_coo和scatter的功能差不多，但它只支持index的shape[0]为1的状态，即每一行都为相同的分组方式。同时，index中数值为顺序排列，以提高计算速度。

torch_scatter.segment_csr

torch_scatter.segment_csr的index格式不太相同，是一种区间格式，如[0, 2, 5]，表示0，1为一组，2，3，4为一组，即取数值间的左闭右开区间。这个方法是计算速度最快的。

官方文档地址

torch_scatter库doc
https://pytorch-scatter.readthedocs.io/en/latest/functions/scatter.html

torch.scatter文档
https://pytorch-cn.readthedocs.io/zh/latest/package_references/Tensor/#scatter_input-dim-index-src-tensor