图机器学习调研洞察：PyG与DGL

图神经网络（Graph Neural Networks，GNN）是人工智能领域最火热的课题之一，近年来，越来越多的研究机构和企业投入大量精力研究图神经网络，推出了自己的图学习引擎，并将其广泛应用在社交网络、电商推荐、欺诈检测、知识图谱等领域。目前，主流的开源图学习引擎有DGL，PyG，GraphScope，Angel，PGL等，其基本信息如下：

可以看出，DGL和PyG支持的算法数量最多，且社区活跃度很高，充分反映了它们在图学习领域的受欢迎程度。

除此之外，一款优秀的图学习引擎，其应该具备强悍的训练和推理性能，针对此，我们基于ogbn-products数据集（点规模2449029，边规模61859140，无向同构图，训练集:验证集:测试集=98:1:1，10个epoch，CPU训练）和GCN算法，对上述图学习引擎进行了性能对比测试，测试结果如下：

综合上述测试结果可以看出，DGL性能最好、内存占用最低，PyG和GraphScope表现次之，Angel和PGL表现最差。

图在化学材料领域可以很好的模拟成原子，而边则是原子间的相互作用，相较于传统的神经网络对于原子位置等有更为自然的定义。DGL与PyG在AI for Science领域被广泛的应用在从小分子到大分子的各个尺度的性质预测和模拟之中。

根据综述中列举的情况，我们在生物制药与材料化学场合进行了一个统计。在AI4Sci领域内，PyG相对DGL的应用更多，在2019年发布之后，基本能够占据40-50%的份额。而在细分领域中，小分子包括晶体等领域使用较多，蛋白质的性质结构预测中应用较少。

具体从实现方式来看，PyG和DGL存在如下不同：

• DGL将每个图转变为DGLGraph，消息传递也是已DGLGraph为中心进行的。通过定义Message Function和Reduce function对DGLGraph中的边和节点进行操作，然后通过DGLgraph的send、recv及updata_all等方法进行更新。

• PyG将每个图储存在一个Data中，但是消息传递是基于MessagePassing基类进行的，与Data或者Batch并无直接联系，通过在网络中重写message passing的forward，message，aggregate和update等方法实现自定义的消息传递过程，Propagete方法会自动调用这些方法完成数据的更新。

值得指出的是，PyG与DGL分别代表了两种不同的图机器学习计算框架。一种是以DGL、PGL为代表的Graph Centirc方式，图上的传播运算和就地运算都在图本身实现和完成，保留了图结构，能支持更为复杂的采样等操作。另一种是以PyG为代表的Tensor Centirc方式（Euler、AliGraph也采用这种方式）：

• 图网络forward时，将图表达为三个Tensor
• edge_index: Tensor, shape:[2, num_edges], 2为src与dst节点ID，num_edges为边个数
• node_feat : Tensor, shape:[num_edges, num_node_feat], 节点特征
• edge_feat : Tensor, shape:[num_edges, num_edge_feat], 边特征
• MessagePassing定义为一种Layer，Message, Reduce等成员函数在Layer中定义

这种方式的缺点是图结构被打开，不够直观，并且图信息被拆成多个Tensor限制了图表达能力。

根据Zhou Y等人做的实验来看，当节点和边的数量较小时，PyG的性能相对更好，而当节点和边的数量较大时，DGL相对PyG有一定的性能优势。

PyG与DGL的架构图如下所示，对于大图性能优势的初步探讨如下。PyG切分成小图之后时串行运行的，而DGL可以并行处理；PyG是基于动态图的，会重复建图，DGL支持动态和静态图，可以避免重复建图的耗时；DGL底层基于CUDA实现，支持PT，TF等AI框架，PyG基于PT，不支持TF等其他框架；PyG底层用到COO，没有用到CSR，DGL两者都支持，之前在CPU上测试结果表明CSR相比COO是有性能提速的。

参考文献

[1] Zhou Y, Song Y, Leng J, et al. AdaptGear: Accelerating GNN Training via Adaptive Subgraph-Level Kernels on GPUs[C]//Proceedings of the 20th ACM International Conference on Computing Frontiers. 2023: 52-62.