图机器学习调研洞察: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.
评论