计图开源：图机器学习库

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

Part 1

问题和背景

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

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

Part 2

Jittor Geometric框架及特性

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

图1 Jittor Geometric的整体架构

Jittor Geometric具有三大特性。

高效灵活

Jittor Geometric对图学习中的各种操作进行了针对性优化设计，框架的高性能主要来自于：

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

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

海量数据集支持

Jittor Geometric支持图机器学习中广泛使用的图数据集(见图2)，其中包括：

引文网络数据集：Cora, CiteSeer, PubMed, Ogbn-Arxiv等
异配图数据集：Chameleon, Squirrel, Penn94,Roman_Empire, Tolokers等
动态图数据集：Wikipedia, Reddit, MOOC, LastFM等
分子图数据集：QM9, ESOL, FreeSolv, BACE等。

图2 Jittor Geometric支持多种图数据集

此外，用户可以通过继承InMemoryDataset类来导入自定义数据集，使用Jittor Geometric内置的partition_graph, to_undirected, remove_isolated_nodes等API来实现图划分，转无向图，去除孤立点等图数据处理操作。

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

Jittor Geometric实现了多个经典的图神经网络模型，还特别支持以下前沿领域的模型[见图3]：

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

图3 Jittor Geometric支持的图神经网络模型

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

深层图卷积模型GCNII[2]: 解决了图卷积网络的过度平滑的问题，并从理论上证明了该算法的有效性。在国际知名机器学习数据网站PaperWithCode的各测评榜单中长期占据第一名的位置，在Cora和Pubmed数据集上的预测准确率自2020年提出至今始终保持第一，获世界人工智能大会青年优秀论文提名奖。自20年来引用数达1793，引用数在近四年ICML会议图神经网络领域的论文中排名第一。
分子图基座模型Uni-Mol[3]:基于分子三维结构的通用分子表征学习框架, 性能优越、模型泛化能力强，在小分子性质预测、蛋白靶点预测和蛋白-配体复合物构象预测等任务上都超越之前方法。

Part 3

未来计划

未来，Jittor Geometric 将持续更新，并计划在以下几个方面进行进一步提升：

1. 更多前沿模型补充

拓展和完善图神经网络领域的最新算法和模型，支持更多的前沿图神经网络架构。

2. 高效动态图计算支持

加强对动态图的高效计算支持，提升图神经网络在处理动态、时序图数据时的性能。

3. 分布式训练加速

进一步优化分布式训练框架，支持多节点、多GPU的并行计算，提升训练大规模图模型的速度和效率。

4. 辅助转换脚本

提供PyTorch Geometric实现到Jittor Geometric实现的自动转换脚本，帮助熟悉PyTorch的用户更快上手。

Jittor Geometric是一个开源项目，欢迎更多研究人员和开发者的参与。如果您有任何问题或建议，或希望贡献代码，请随时联系我们。

Email:

runlin_lei@ruc.edu.cn

开源地址：

https://github.com/AlgRUC/JittorGeometric

感谢所有参与和支持Jittor Geometric项目的开发者和研究者！感谢清华大学计图团队的技术支持！

参考文献

Wang Q, Zhang Y, Wang H, et al., Neutronstar: distributed GNN training with hybrid dependency management, International Conference on Management of Data, 2022, 1301-1315.
Chen M, Wei Z, Huang Z, et al., Simple and deep graph convolutional networks, International Conference on Machine Learning, 2020, article No. 161, 1725-1735.
Zhou G, Gao Z, Ding Q, et al., Uni-mol: A universal 3d molecular representation learning framework, International Conference on Learning Representations, 2023.

GGC往期回顾

1. 计图适配华为昇腾910系列，迈向国产AI新生态

2. 张量场引导的机器人分层路径规划与自主探索 | CVMJ Spotlight

3. 跨模态增强的长尾图像分类 | CVMJ Spotlight

4. 计图开源：基于扩散模型的语音驱动风格化人脸动画生成

5. 第四届计图人工智能算法挑战赛闭幕，16支队伍获奖！

可通过下方二维码，关注清华大学图形学实验室，了解图形学、深度学习、Jittor框架、CVMJ期刊和CVM会议的相关资讯。

计图开源：图机器学习库