近日，中国人民大学与东北大学联合开发了图机器学习库Jittor Geometric，其1.0版本近日已正式发布。

Jittor Geometric以国产深度学习框架Jittor为基础技术架构，聚焦图数据，在图存储、图计算、图学习等方面作了细致优化，整合、加速了现有多类图神经网络模型，模型运行时间在多种图学习任务上较Pytorch Geometric（PyG）、Deep Graph Library（DGL）等同类型框架提升10%~50%。

同时，Jittor Geometric简洁、易用、跨平台通用性强、用户学习成本低，目前已用于研究生课程教学。

开源地址：

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

图 1：Jittor Geometric架构示意图

目前，Jittor Geometric已集成了以谱域、动态、几何图神经网络模型为代表的多类代表性图学习模型。

Jittor Geometric的底层为国产深度学习框架Jittor，在其上包含图数据存储层、图计算优化层和图学习设计层。

其中，图存储层采用了压缩存储、内存调度等策略以减少显存消耗；图计算层通过图算子设计、图结构划分、并行计算及稀疏矩阵高效计算等方法加速图计算效率。

图学习层将图神经网络的传播学习过程拆分为几类代表性图算子的组合，进而通过加速图算子的数据存取与计算效率加速图学习模型的性能表现。

据研发团队介绍， Jittor Geometric具有三大特性：

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

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

海量图数据集支持

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

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

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

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

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

图 3：Jittor Geometric支持多领域图神经网络模型

其支持多种前沿模型：

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将在以下几个方面进行进一步提升：