信息来源：中国人民大学网站 发布时间  2025年09月06日

在国家自然科学基金重点项目“面向图学习的分布式原生图存储和计算方法研究”的支持下，中国人民大学魏哲巍教授团队与东北大学张岩峰教授团队联合发布了基于计图(Jittor)深度学习框架的图机器学习库(Jittor Geometric)。该库对图数据的存储和训练作了针对性的优化，支持多个图学习前沿领域的图神经网络模型与数据集，为图机器学习的研究提供了更高效、更关注前沿进展的新框架。

图数据在各类任务中的应用越来越广泛，已成为多领域研究的核心。无论是在社交网络中的用户与关系、推荐系统中的用户与物品交互，还是生物信息学中的蛋白质相互作用、金融风控中的交易网络，图结构数据的价值都无可替代。图神经网络（Graph Neural Networks, GNNs）作为一种强大的图数据处理工具，利用消息传递机制，能够在图结构数据中高效地学习节点、边以及它们之间的关系，从而得到深层次的节点表示。与传统的神经网络方法相比，图神经网络可以直接处理图数据，有效提升模型对图数据的理解能力和应用效果。

JittorGeometric是基于国内自主研发的深度学习框架 Jittor打造的图机器学习库，旨在为图神经网络的研究者和开发者提供一个高效、灵活的工具。Jittor使用即时编译和元算子技术，框架具有高效性、灵活性和可定制性。作为第一个支持Jittor框架的国产图机器学习平台，JittorGeometric融合了Jittor框架的核心特性，为图神经网络的研究和应用提供显著的性能提升。特别是在大规模图数据的处理和训练中，JittorGeometric展现出了优越的计算效率和高效的资源利用率，助力研究人员高效运行图学习任务。

（图1 JittorGeometric的整体架构）

JittorGeometric的整体架构如图1所示。在底层，基于Jittor深度学习框架针对图数据的存储和训练进行优化：通过压缩存储与内存调度等减少显存消耗，优化图存储；通过图算子、图划分、并行计算、 warp级调度及稀疏矩阵计算优化等[1]加速图计算，大幅减少训练时间。在顶层，由各种基础图算子组成的JittorGeometric API涵盖了图采样、消息传递聚合、图数据读取处理和多种图卷积层。基于API，用户可以自由设计并实现各种领域的图神经网络模型，包括前沿的谱域、动态、几何图神经网络模型。

JittorGeometric具有以下特性：

    （1）高效灵活：JittorGeometric对图学习中的各种操作进行了针对性优化设计，框架的高性能主要来自于：
 

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

基于JittorGeometric实现的图神经网络模型拥有更快的训练速度。在单块NVIDIA A100显卡上进行测试，相较Pytorch Geometric和Deep Graph Learning等主流框架在Cora, Citeseer, Pubmed经典图数据集上提升25%。

（图2 JittorGeometric 相较现有图机器学习库的性能优势  ）

    （2）海量数据集支持
 

JittorGeometric支持图机器学习中广泛使用的图数据集，其中包括：引文网络数据集、异配图数据集、动态图数据集、分子图数据集等。JittorGeometric 同时提供了相对应的数据集预处理接口，旨在为广泛的图机器学习应用提供高效、可靠的数据支持。

（图3 支持海量图数据集）

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

JittorGeometric实现了多个经典的图神经网络模型，还特别支持以下前沿领域的模型：

•        动态图神经网络模型: 增强对动态图的支持，提升动态图场景下的模型性能。
•       谱域图神经网络模型: 包含谱图神经网络方法，促进谱图学习在不同领域的应用。
•       分子图神经网络模型: 支持前沿分子图神经网络模型，推动生物领域研究。
•       图表示学习：支持经典图表示学习算法，包含无监督场景下经典算法。
•       图推荐系统算法：提供在推荐系统中得到广泛应用的图学习算法实现，助力推荐系统领域研究。

（图4 JittorGeometric包含的图神经网络模型）

其中包括许多著名模型的实现：

•       深层图卷积模型GCNII[2]: 解决了图卷积网络的过度平滑的问题，并从理论上证明了该算法的有效性。在国际知名机器学习数据网站  PaperWithCode的各测评榜单中长期占据第一名的位置，在Cora和Pubmed数据集上的预测准确率自2020年提出至今始终保持第一，获世界人工智能大会青年优秀论文提名奖。自20年来引用数达2181，引用数在近五年ICML会议图神经网络领域的论文中排名第一。
•       分子图基座模型Uni-Mol[3]: 基于分子三维结构的通用分子表征学习框架, 性能优越、模型泛化能力强，在小分子性质预测、蛋白靶点预测和蛋白-配体复合物构象预测等任务上都超越之前方法。
•       推荐领域经典模型LightGCN[4]: 图卷积网络在推荐系统中的应用里程碑。该模型因其卓越的有效性、极强的简洁性和高效性，成为推荐系统领域最广泛采用的基线模型之一，在多个公开基准数据集上实现了最先进的性能。谷歌学术引用数达4900+。

（4）并行计算优化

Minibatch 训练支持：JittorGeometric引入了基于稀疏图格式的Minibatch训练范式，使得在大规模图数据上进行批量训练成为可能，显著降低了显存开销，并提高了计算效率。在GPU场景下，训练效率超过现有图机器学习框架。

分布式多卡计算：JittorGeometric在分布式计算架构上做了重要扩展，支持数据并行和张量并行两种训练模式，能够多卡进行大规模图数据训练，极大地提升了计算速度。

（图5 JittorGeometric 提供高效的数据并行（上）与张量并行计算（下））

（5）国产硬件支持

国产硬件支持：JittorGeometric针对昇腾NPU进行了优化，利用核心间数据并行和核心内流水行并行策略，加速了图神经网络的训练过程，特别在大规模图数据的处理上提升了性能。随着国产硬件的逐步发展，JittorGeometric在处理图神经网络训练时能够充分发挥国产芯片的性能，提供硬件加速支持。

（图6 JittorGeometric 提供国产芯片上的高效实现）

在未来，JittorGeoemetric将持续维护，计划在以下方面进一步提升：

     一是更多前沿模型补充。 就图机器学习更多方面，如异构图、图基座模型等领域进行模型补充，并提供相应模型适配。对已包含场景继续丰富集成算法。
      二是持续提升计算效率。在分布式场景继续优化性能，尝试多节点进行大规模图训练。另外，在异构集群（机器算力和通信异构）的情况下执行高效的任务分配与模型训练。
      三是国产硬件性能优化。在国产硬件（昇腾）上，进一步扩大支持训练集的规模，探索多卡训练和分布式训练。

JittorGeometric是一个开源项目，欢迎更多研究人员和开发者的参与。
开源地址：https://github.com/AlgRUC/JittorGeometric