近日，中國人民大學與東北大學聯(lián)合開發(fā)了圖機器學習庫Jittor Geometric，其1.0版本近日已正式發(fā)布。

Jittor Geometric以國產深度學習框架Jittor為基礎技術架構，聚焦圖數據，在圖存儲、圖計算、圖學習等方面作了細致優(yōu)化，整合、加速了現有多類圖神經網絡模型，模型運行時間在多種圖學習任務上較Pytorch Geometric（PyG）、Deep Graph Library（DGL）等同類型框架提升10%~50%。

同時，Jittor Geometric簡潔、易用、跨平臺通用性強、用戶學習成本低，目前已用于研究生課程教學。

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

「Jittor Geometric」架構全景圖及三大特性

圖 1：Jittor Geometric架構示意圖

目前，Jittor Geometric已集成了以譜域、動態(tài)、幾何圖神經網絡模型為代表的多類代表性圖學習模型。

Jittor Geometric的底層為國產深度學習框架Jittor，在其上包含圖數據存儲層、圖計算優(yōu)化層和圖學習設計層。

其中，圖存儲層采用了壓縮存儲、內存調度等策略以減少顯存消耗；圖計算層通過圖算子設計、圖結構劃分、并行計算及稀疏矩陣高效計算等方法加速圖計算效率。

圖學習層將圖神經網絡的傳播學習過程拆分為幾類代表性圖算子的組合，進而通過加速圖算子的數據存取與計算效率加速圖學習模型的性能表現。

據研發(fā)團隊介紹， Jittor Geometric具有三大特性：

高效性與靈活性的完美結合

在編寫Jittor Geometric之初，研發(fā)團隊就對圖學習中的各種操作進行了針對性優(yōu)化設計，使框架兼顧高效性與靈活性。研發(fā)團隊表示Jittor Geometric的高性能主要來自于以下五個方面：

• AVX向量化: 通過AVX指令集進行向量化計算，顯著提高了大規(guī)模數據處理和計算效率。
• 頂點級別并行計算: 利用頂點級并行計算技術，有效提升了圖結構數據的處理速度。
• 稀疏矩陣算子設計：采用稀疏矩陣乘法算子（SPMM），大幅提升稀疏圖傳播過程計算效率。
• 基于warp的并行計算: 在GPU中使用warp級別的并行計算，實現高效圖數據處理。
• CPU-GPU自動內存交換: 實現CPU與GPU之間的自動內存管理與數據交換，優(yōu)化內存使用和數據傳輸效率。

海量圖數據集支持

圖 2：Jittor Geometric支持海量圖數據集

Jittor Geometric支持多種類型圖數據，涵蓋了圖神經網絡各領域經典任務的數據集，用戶可以用統(tǒng)一方式來高效便捷地讀取和調用。

此外，研發(fā)團隊還提供了靈活的API，便于用戶導入自定義數據集進行實驗。

多領域圖機器學習模型支持

Jittor Geometric不僅實現了經典的圖神經網絡模型，還特別涵蓋圖神經網絡前沿領域的模型：

• 動態(tài)圖神經網絡模型：增強對動態(tài)圖的支持，提升動態(tài)圖場景下的模型性能。
• 譜域圖神經網絡模型：包含譜圖神經網絡方法，促進譜圖學習在不同領域的應用。
• 分子圖神經網絡模型：支持前沿分子圖神經網絡模型，推動生物領域研究。

圖 3：Jittor Geometric支持多領域圖神經網絡模型

其支持多種前沿模型：

• 深層圖卷積模型GCNII: 解決了圖卷積網絡的過度平滑的問題，并從理論上證明了該算法的有效性。在國際知名機器學習數據網站PaperWithCode的各測評榜單中長期占據第一名的位置，在Cora和Pubmed數據集上的預測準確率自2020年提出至今始終保持第一，獲世界人工智能大會青年優(yōu)秀論文提名獎。
• 分子圖基座模型Uni-Mol: 基于分子三維結構的通用分子表征學習框架, 性能優(yōu)越、模型泛化能力強，在小分子性質預測、蛋白靶點預測和蛋白-配體復合物構象預測等任務上都超越之前方法。

Jittor Geometric的各種模型均以高度統(tǒng)一的模式編寫，極大降低了用戶的學習成本。在熟練掌握一類數據集和模型后，用戶可以借助Jittor Geometric自由探索各領域圖神經網絡模型的設計，無需花費大量時間進行領域間的遷移學習。

據研發(fā)團隊透露，未來Jittor Geometric還將進一步拓展支持的數據集和模型，關注圖神經網絡前沿研究，如大語言模型與圖學習等，助力研究人員和開發(fā)者們更便捷地進行前沿探索和落地應用。

多個數據集上，性能超過現有主流框架

Jittor Geometric開發(fā)團隊提供了實驗數據。以GCN模型為例，在各種經典圖數據集上Jittor Geometric均展現出與現有主流圖機器學習框架相當或更優(yōu)的性能。

具體來說，Jittor Geometric在小規(guī)模圖數據集上的訓練速度相較主流框架提升25%；在大規(guī)模圖數據集上顯著優(yōu)于Pytorch Geometric和Paddle Graph Learning, 與Deep Graph Library相當。

圖 4：Jittor Geometric與各主流框架性能對比

為了方便廣大用戶上手Jittor Geometric，研發(fā)團隊采用了與PyTorch Geometric較為相似的模塊化接口。從上手難度來看，熟悉 PyTorch Geometric的用戶可以直接上手，甚至沒有因為變量名不同而帶來影響。

Jittor Geometric提供了良好的課程支持，幫助對圖機器學習感興趣的同學們更好地學習和設計各種圖神經網絡。一方面，Jittor Geometric 定義了全面且精簡的基礎圖算子，同學們可以基于這些基礎算子實現多樣的消息傳遞和圖算法；另一方面還提供了各領域的模型實現、訓練框架及數據集，幫助同學們更快上手圖機器學習領域的任務、數據及模型設計。

圖 5：Jittor Geometric在圖機器學習課程中的應用

據研發(fā)團隊介紹，Jittor Geometric已經應用到課程教學中，同學們基于該框架設計并實現了各種動態(tài)圖神經網絡和異配圖神經網絡，在多個數據集上進行實驗并作相關分析。同學們反饋該框架簡明易懂，學習成本低，在多種前沿模型及數據集上相較Torch Geometric實現更簡單，訓練更快速。

Jittor Geometric的新進展

Jittor Geometric作為一個新興圖機器學習平臺，在一些功能上，仍舊需要持續(xù)迭代完善。比如生態(tài)的建設，以及更大范圍的推廣，仍舊需要很多的努力。

據研發(fā)團隊透露，Jittor Geometric將在以下幾個方面進行進一步提升：

1. 更多前沿模型補充

拓展和完善圖神經網絡領域的最新算法和模型，支持更多的前沿圖神經網絡架構。

2. 高效動態(tài)圖計算支持

加強對動態(tài)圖的高效計算支持，提升圖神經網絡在處理動態(tài)、時序圖數據時的性能。

3. 分布式訓練加速

進一步優(yōu)化分布式訓練框架，支持多節(jié)點、多GPU的并行計算，提升訓練大規(guī)模圖模型的速度和效率。

4. 輔助轉換腳本

提供PyTorch Geometric實現到Jittor Geometric實現的自動轉換腳本，幫助熟悉PyTorch的用戶更快上手。

研究團隊: 中國高校圖學習開源框架的新力量

據官網介紹，Jittor Geometric由中國人民大學魏哲巍教授和東北大學張巖峰教授及其團隊共同打造。

開發(fā)Jittor Geometric的主力，是實驗室雷潤林，陸施展，付振波等一批博士生。他們的主要研究方向涵蓋譜域、動態(tài)、分子等多領域圖神經網絡，圖存儲與圖計算優(yōu)化等，在重要國際會議NeurIPS, ICML, ICLR, KDD, VLDB, SIGMOD, ICDE等發(fā)表多篇論文。據雷潤林透露，團隊接下來的重點將圍繞學術界展開，致力于使Jittor Geometric成為國內外學術界最受歡迎、使用最廣泛的圖機器學習庫之一，并對AI產業(yè)界產生積極影響。

雷潤林表示：「我們聚焦于圖機器學習的前沿領域，關注圖學習未來發(fā)展中關鍵的動態(tài)圖和大規(guī)模圖數據處理等。同時希望能夠通過構建活躍的開源社區(qū)，不僅完善框架本身，還形成各種圖學習的實踐指引，技術博客及問題討論等，幫助對圖機器學習領域感興趣的用戶更好地探索實踐。目前，Jittor Geometric已經應用到課堂教學中并取得良好反饋，未來還將聯(lián)合更多高校將其融入教學體系中。」

總之，打造一款全面高效的圖機器學習庫并非易事，而選擇開源則意味著要接受來自全球開發(fā)者社區(qū)的嚴格檢驗。

如今，Jittor Geometric已經邁出了至關重要的一步，誠邀所有感興趣的讀者參與其中，一同見證這一新興框架的成長與發(fā)展。