人大、东北大学联合开发「图机器学习库」Jittor Geometric！性能超越PYG、DGL

作者


  新智元报道  

编辑：LRST
【新智元导读】Jittor Geometric 1.0是由中国人民大学与东北大学联合开发的图机器学习库，基于国产Jittor框架，高效灵活，可助力处理复杂图结构数据，性能优于同类型框架，支持多种前沿图神经网络模型，已开源供用户使用。


近日，中国人民大学与东北大学联合开发了图机器学习库Jittor Geometric，其1.0版本近日已正式发布。
Jittor Geometric以国产深度学习框架Jittor为基础技术架构，聚焦图数据，在图存储、图计算、图学习等方面作了细致优化，整合、加速了现有多类图神经网络模型，模型运行时间在多种图学习任务上较Pytorch Geometric（PyG）、Deep Graph Library（DGL）等同类型框架提升10%~50%。
同时，Jittor Geometric简洁、易用、跨平台通用性强、用户学习成本低，目前已用于研究生课程教学。

开源地址：https://github.com/AlgRUC/JittorGeometric

「Jittor Geometric」架构全景图及三大特性

图 1：Jittor Geometric架构示意图

目前，Jittor Geometric已集成了以谱域、动态、几何图神经网络模型为代表的多类代表性图学习模型。
Jittor Geometric的底层为国产深度学习框架Jittor，在其上包含图数据存储层、图计算优化层和图学习设计层。
其中，图存储层采用了压缩存储、内存调度等策略以减少显存消耗；图计算层通过图算子设计、图结构划分、并行计算及稀疏矩阵高效计算等方法加速图计算效率。
图学习层将图神经网络的传播学习过程拆分为几类代表性图算子的组合，进而通过加速图算子的数据存取与计算效率加速图学习模型的性能表现。
据研发团队介绍， Jittor Geometric具有三大特性：

高效性与灵活性的完美结合

在编写Jittor Geometric之初，研发团队就对图学习中的各种操作进行了针对性优化设计，使框架兼顾高效性与灵活性。研发团队表示Jittor Geometric的高性能主要来自于以下五个方面：

  • AVX向量化: 通过AVX指令集进行向量化计算，显著提高了大规模数据处理和计算效率。
  • 顶点级别并行计算: 利用顶点级并行计算技术，有效提升了图结构数据的处理速度。
  • 稀疏矩阵算子设计：采用稀疏矩阵乘法算子（SPMM），大幅提升稀疏图传播过程计算效率。
  • 基于warp的并行计算: 在GPU中使用warp级别的并行计算，实现高效图数据处理。
  • CPU-GPU自动内存交换: 实现CPU与GPU之间的自动内存管理与数据交换，优化内存使用和数据传输效率。

海量图数据集支持

图 2：Jittor Geometric支持海量图数据集

Jittor Geometric支持多种类型图数据，涵盖了图神经网络各领域经典任务的数据集，用户可以用统一方式来高效便捷地读取和调用。

此外，研发团队还提供了灵活的API，便于用户导入自定义数据集进行实验。

多领域图机器学习模型支持

Jittor Geometric不仅实现了经典的图神经网络模型，还特别涵盖图神经网络前沿领域的模型：

  • 动态图神经网络模型：增强对动态图的支持，提升动态图场景下的模型性能。
  • 谱域图神经网络模型：包含谱图神经网络方法，促进谱图学习在不同领域的应用。
  • 分子图神经网络模型：支持前沿分子图神经网络模型，推动生物领域研究。
图 3：Jittor Geometric支持多领域图神经网络模型

其支持多种前沿模型：

  • 深层图卷积模型GCNII: 解决了图卷积网络的过度平滑的问题，并从理论上证明了该算法的有效性。在国际知名机器学习数据网站PaperWithCode的各测评榜单中长期占据第一名的位置，在Cora和Pubmed数据集上的预测准确率自2020年提出至今始终保持第一，获世界人工智能大会青年优秀论文提名奖。

  • 分子图基座模型Uni-Mol: 基于分子三维结构的通用分子表征学习框架, 性能优越、模型泛化能力强，在小分子性质预测、蛋白靶点预测和蛋白-配体复合物构象预测等任务上都超越之前方法。

Jittor Geometric的各种模型均以高度统一的模式编写，极大降低了用户的学习成本。在熟练掌握一类数据集和模型后，用户可以借助Jittor Geometric自由探索各领域图神经网络模型的设计，无需花费大量时间进行领域间的迁移学习。

据研发团队透露，未来Jittor Geometric还将进一步拓展支持的数据集和模型，关注图神经网络前沿研究，如大语言模型与图学习等，助力研究人员和开发者们更便捷地进行前沿探索和落地应用。

多个数据集上，性能超过现有主流框架

Jittor Geometric开发团队提供了实验数据。以GCN模型为例，在各种经典图数据集上Jittor Geometric均展现出与现有主流图机器学习框架相当或更优的性能。

具体来说，Jittor Geometric在小规模图数据集上的训练速度相较主流框架提升25%；在大规模图数据集上显著优于Pytorch Geometric和Paddle Graph Learning, 与Deep Graph Library相当。

图 4：Jittor Geometric与各主流框架性能对比

为了方便广大用户上手Jittor Geometric，研发团队采用了与PyTorch Geometric较为相似的模块化接口。从上手难度来看，熟悉 PyTorch Geometric的用户可以直接上手，甚至没有因为变量名不同而带来影响。

Jittor Geometric提供了良好的课程支持，帮助对图机器学习感兴趣的同学们更好地学习和设计各种图神经网络。一方面，Jittor Geometric 定义了全面且精简的基础图算子，同学们可以基于这些基础算子实现多样的消息传递和图算法；另一方面还提供了各领域的模型实现、训练框架及数据集，帮助同学们更快上手图机器学习领域的任务、数据及模型设计。

图 5：Jittor Geometric在图机器学习课程中的应用

据研发团队介绍，Jittor Geometric已经应用到课程教学中，同学们基于该框架设计并实现了各种动态图神经网络和异配图神经网络，在多个数据集上进行实验并作相关分析。同学们反馈该框架简明易懂，学习成本低，在多种前沿模型及数据集上相较Torch Geometric实现更简单，训练更快速。

Jittor Geometric的新进展

Jittor Geometric作为一个新兴图机器学习平台，在一些功能上，仍旧需要持续迭代完善。比如生态的建设，以及更大范围的推广，仍旧需要很多的努力。

据研发团队透露，Jittor Geometric将在以下几个方面进行进一步提升：

1. 更多前沿模型补充
拓展和完善图神经网络领域的最新算法和模型，支持更多的前沿图神经网络架构。
2. 高效动态图计算支持
加强对动态图的高效计算支持，提升图神经网络在处理动态、时序图数据时的性能。
3. 分布式训练加速
进一步优化分布式训练框架，支持多节点、多GPU的并行计算，提升训练大规模图模型的速度和效率。
4. 辅助转换脚本
提供PyTorch Geometric实现到Jittor Geometric实现的自动转换脚本，帮助熟悉PyTorch的用户更快上手。

研究团队: 中国高校图学习开源框架的新力量

据官网介绍，Jittor Geometric由中国人民大学魏哲巍教授和东北大学张岩峰教授及其团队共同打造。

开发Jittor Geometric的主力，是实验室雷润林，陆施展，付振波等一批博士生。他们的主要研究方向涵盖谱域、动态、分子等多领域图神经网络，图存储与图计算优化等，在重要国际会议NeurIPS, ICML, ICLR, KDD, VLDB, SIGMOD, ICDE等发表多篇论文。据雷润林透露，团队接下来的重点将围绕学术界展开，致力于使Jittor Geometric成为国内外学术界最受欢迎、使用最广泛的图机器学习库之一，并对AI产业界产生积极影响。

雷润林表示：「我们聚焦于图机器学习的前沿领域，关注图学习未来发展中关键的动态图和大规模图数据处理等。同时希望能够通过构建活跃的开源社区，不仅完善框架本身，还形成各种图学习的实践指引，技术博客及问题讨论等，帮助对图机器学习领域感兴趣的用户更好地探索实践。目前，Jittor Geometric已经应用到课堂教学中并取得良好反馈，未来还将联合更多高校将其融入教学体系中。」

总之，打造一款全面高效的图机器学习库并非易事，而选择开源则意味着要接受来自全球开发者社区的严格检验。

如今，Jittor Geometric已经迈出了至关重要的一步，诚邀所有感兴趣的读者参与其中，一同见证这一新兴框架的成长与发展。

（文：新智元）

欢迎分享

发表评论