Equitorch: 基于pyg的模块化等变图神经网络包

文章预览

中山大学陈川课题组（GTML）近日发布了等变图神经网络包Equitorch。Equitorch包中以模块化的形式集成了大多等变算子，并基于pyg构建图神经网络，同时提供了详细的文档、示例与教程。 背景 近两年来，随着AI for Science任务的兴起，3D等变图神经网络（3D-Equivariant GNN）获得了广泛的关注，出现了eSCN、NeuqIP、SEGNN、Equiformer(V2)、DPA(-2)等模型，并在分子几何生成、能量预测、粒子运动求解等领域获得了优异表现。                 图片说明：3D旋转等变函数：φ(Rx)=Rφ(x)，若将函数的输入进行旋转，则函数的输出也会相应旋转。(图片源自https://arxiv.org/abs/2110.02905) 特点 在Equitorch中，我们以一种模块化的方式集成了现在工作中使用的大多等变算子以及其它相关辅助函数，并统一这些操作的数据格式约定。 对图神经网络相关的部分，如消息传递机制、图数据 ………………………………