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

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

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

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

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