PyG教程(7)：剖析邻域聚合

一.前言

上篇文章《PyG教程(6)：自定义消息传递网络》主要介绍了消息传递GNN的大致框架。本文主要聚焦于消息传播中的邻域聚合，本文将介绍PyG是如何将节点的邻居的消息聚合到节点本身的。

二.PyG中的邻域聚合

PyG中邻域聚合是通过aggregate(inputs, index)函数来完成的，该函数的第一个参数inputs为消息构建函数message()构建的消息，该函数还存在一个参数index，这个参数对于消息聚合是十分关键的，它指示了inputs中每条消息属于哪个节点的邻域。下图便很好的解释了PyG中的消息聚合：

上述栗子中展示的是包含4个顶点、8条边的graph，其中input为在8条边上传播的消息、index为各条边上消息的归属，即目标节点的索引。通过index，可以将属于同一个节点邻域的消息聚合到一起，常见的聚合包括sum、mean、mean、mul和min等。

在PyG中通过scatter函数来实现上述过程，查看MessagePassing的源码，可以看到其aggregate函数的定义如下：

def aggregate(self, inputs: Tensor, index: Tensor,
              ptr: Optional[Tensor] = None,
              dim_size: Optional[int] = None) -> Tensor:
    r""" 注释太长略 """
    if ptr is not None:
        ptr = expand_left(ptr, dim=self.node_dim, dims=inputs.dim())
        return segment_csr(inputs, ptr, reduce=self.aggr)
    else:
        return scatter(inputs, index, dim=self.node_dim, dim_size=dim_size,
                       reduce=self.aggr)

aggregate函数中scatter函数源码为：

def scatter(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
            out: Optional[torch.Tensor] = None, dim_size: Optional[int] = None,
            reduce: str = "sum") -> torch.Tensor:
    r""" 注释太长略 """
    if reduce == 'sum' or reduce == 'add':
        return scatter_sum(src, index, dim, out, dim_size)
    if reduce == 'mul':
        return scatter_mul(src, index, dim, out, dim_size)
    elif reduce == 'mean':
        return scatter_mean(src, index, dim, out, dim_size)
    elif reduce == 'min':
        return scatter_min(src, index, dim, out, dim_size)[0]
    elif reduce == 'max':
        return scatter_max(src, index, dim, out, dim_size)[0]
    else:
        raise ValueError

其中便包含了前面提到的5种聚合方式。对于这些聚合方式，只需要在继承MessagePassing类时，通过super().__init__来向该类传递参数aggr参数的值即可。

三.torch_scatter模块

若用户需要自定义消息聚合，则在重写的aggregate()函数中，同样可以使用MessagePassing中的scatter函数，只需要导入torch_scatter模块即可。

在torch_scatter模块中也实现了scatter函数，其声明如下：

scatter: (src: Tensor, index: Tensor, dim: int = -1, out: Tensor | None = None, dim_size: int | None = None, reduce: str = "sum") -> Tensor

常用参数说明：

参数	说明
src	每条边上的源节点生成的消息
index	指示每条边上消息需要聚合到哪个节点上
dim	指示沿着那个维度（轴）应用index进行聚合
reduce	聚合操作，包括sum、mul、 mean、 min 和 max

注意，torch_scatter也为上述的几种聚合单独提供了API：

torch_scatter.scatter_add()
torch_scatter.scatter_max()
torch_scatter.scatter_mean()
torch_scatter.scatter_min()
torch_scatter.scatter_mul()

为了方便理解，下面给出一个栗子，假设存在一个包含3个顶点、6条边的图：

假设0、1、2三个顶点生成的消息分别为1、2、3，则图中6条边的消息inputs和相应的index构造如下：

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

应用torch_scatter.scatter()函数的结果如下：

out = torch_scatter.scatter(src=inputs, index=index, dim=0, reduce="sum")
print(out)
""" tensor([[5], [4], [3]]) """

可以看到节点0接受来自节点1，2的消息得到2+3=5，节点1接受来自节点0，2的消息得到1+3=4，而节点2接受来自节点0，1的消息得到1+2=3。

四.结语

参考资料：

• torch_scatter

通过本文可以加深对PyG中消息聚合过程的理解，这将有助于更好的自定义GNN模型。以上便是本文的全部内容，若有任何错误，请批评指正。