昇腾AI4S图机器学习：DGL消息传递接口的PyG替换

背景介绍

DGL (Deep Graph Learning) 和 PyG (Pytorch Geometric) 是两个主流的图神经网络库，它们在API设计和底层实现上有一定差异，在不同场景下，研究人员会使用不同的依赖库，昇腾NPU对PyG图机器学习库的支持亲和度更高，因此有些时候需要做DGL接口的PyG替换。

 

SE3Transformer在RFdiffusion蛋白质设计模型中（GitHub - RosettaCommons/RFdiffusion: Code for running RFdiffusion）作为核心组件，负责处理蛋白质结构的几何信息。其架构基于图神经网络，通过SE(3)等变性实现对三维旋转和平移的不变性特征提取。本系列以RFDiffusion模型中的SE3Transformer为例，讲解如何将DGL中的接口替换为PyG实现。

 

在本文中，主要展示消息传递接口的PyG替换。

消息传递接口

一、边-节点消息传递 (EdgeSoftmax + Aggregation)