新智元報道
編輯：LRST【新智元導(dǎo)讀】Jittor Geometric 1.0是由中國人民大學(xué)與東大學(xué)聯(lián)合開發(fā)的圖機(jī)器學(xué)習(xí)庫，基于國產(chǎn)Jittor框架，高效靈活，可助力處理復(fù)雜圖結(jié)構(gòu)數(shù)據(jù)，性能優(yōu)于同類型框架，支持多種前沿圖神經(jīng)網(wǎng)絡(luò)模型，已開源供用戶使用。
近日，中國人民大學(xué)與東北大學(xué)聯(lián)合開發(fā)了圖機(jī)器學(xué)習(xí)庫Jittor Geometric，其1.0版本近日已正式發(fā)布。Jittor Geometric以國產(chǎn)深度學(xué)習(xí)框架Jittor為基礎(chǔ)技術(shù)架構(gòu)，聚焦圖數(shù)據(jù)，在圖存儲、圖計算、圖學(xué)習(xí)等方面作了細(xì)致優(yōu)化，整合、加速了現(xiàn)有多類圖神經(jīng)網(wǎng)絡(luò)模型，模型運(yùn)行時間在多種圖學(xué)習(xí)任務(wù)上較Pytorch Geometric（PyG）、Deep Graph Library（DGL）等同類型框架提升10%~50%。同時，Jittor Geometric簡潔、易用、跨平臺通用性強(qiáng)、用戶學(xué)習(xí)成本低，目前已用于研究生課程教學(xué)。開源地址：https://github.com/AlgRUC/JittorGeometric
「Jittor Geometric」架構(gòu)全景圖及三大特性
圖 1：Jittor Geometric架構(gòu)示意圖
目前，Jittor Geometric已集成了以譜域、動態(tài)、幾何圖神經(jīng)網(wǎng)絡(luò)模型為代表的多類代表性圖學(xué)習(xí)模型。Jittor Geometric的底層為國產(chǎn)深度學(xué)習(xí)框架Jittor，在其上包含圖數(shù)據(jù)存儲層、圖計算優(yōu)化層和圖學(xué)習(xí)設(shè)計層。其中，圖存儲層采用了壓縮存儲、內(nèi)存調(diào)度等策略以減少顯存消耗；圖計算層通過圖算子設(shè)計、圖結(jié)構(gòu)劃分、并行計算及稀疏矩陣高效計算等方法加速圖計算效率。圖學(xué)習(xí)層將圖神經(jīng)網(wǎng)絡(luò)的傳播學(xué)習(xí)過程拆分為幾類代表性圖算子的組合，進(jìn)而通過加速圖算子的數(shù)據(jù)存取與計算效率加速圖學(xué)習(xí)模型的性能表現(xiàn)。據(jù)研發(fā)團(tuán)隊(duì)介紹， Jittor Geometric具有三大特性：高效性與靈活性的完美結(jié)合在編寫Jittor Geometric之初，研發(fā)團(tuán)隊(duì)就對圖學(xué)習(xí)中的各種操作進(jìn)行了針對性優(yōu)化設(shè)計，使框架兼顧高效性與靈活性。研發(fā)團(tuán)隊(duì)表示Jittor Geometric的高性能主要來自于以下五個方面：
AVX向量化: 通過AVX指令集進(jìn)行向量化計算，顯著提高了大規(guī)模數(shù)據(jù)處理和計算效率。
頂點(diǎn)級別并行計算: 利用頂點(diǎn)級并行計算技術(shù)，有效提升了圖結(jié)構(gòu)數(shù)據(jù)的處理速度。
稀疏矩陣算子設(shè)計：采用稀疏矩陣乘法算子（SPMM），大幅提升稀疏圖傳播過程計算效率。
基于warp的并行計算: 在GPU中使用warp級別的并行計算，實(shí)現(xiàn)高效圖數(shù)據(jù)處理。
CPU-GPU自動內(nèi)存交換: 實(shí)現(xiàn)CPU與GPU之間的自動內(nèi)存管理與數(shù)據(jù)交換，優(yōu)化內(nèi)存使用和數(shù)據(jù)傳輸效率。
海量圖數(shù)據(jù)集支持
圖 2：Jittor Geometric支持海量圖數(shù)據(jù)集
Jittor Geometric支持多種類型圖數(shù)據(jù)，涵蓋了圖神經(jīng)網(wǎng)絡(luò)各領(lǐng)域經(jīng)典任務(wù)的數(shù)據(jù)集，用戶可以用統(tǒng)一方式來高效便捷地讀取和調(diào)用。
此外，研發(fā)團(tuán)隊(duì)還提供了靈活的API，便于用戶導(dǎo)入自定義數(shù)據(jù)集進(jìn)行實(shí)驗(yàn)。
多領(lǐng)域圖機(jī)器學(xué)習(xí)模型支持Jittor Geometric不僅實(shí)現(xiàn)了經(jīng)典的圖神經(jīng)網(wǎng)絡(luò)模型，還特別涵蓋圖神經(jīng)網(wǎng)絡(luò)前沿領(lǐng)域的模型：
動態(tài)圖神經(jīng)網(wǎng)絡(luò)模型：增強(qiáng)對動態(tài)圖的支持，提升動態(tài)圖場景下的模型性能。
譜域圖神經(jīng)網(wǎng)絡(luò)模型：包含譜圖神經(jīng)網(wǎng)絡(luò)方法，促進(jìn)譜圖學(xué)習(xí)在不同領(lǐng)域的應(yīng)用。
分子圖神經(jīng)網(wǎng)絡(luò)模型：支持前沿分子圖神經(jīng)網(wǎng)絡(luò)模型，推動生物領(lǐng)域研究。
圖 3：Jittor Geometric支持多領(lǐng)域圖神經(jīng)網(wǎng)絡(luò)模型其支持多種前沿模型：
深層圖卷積模型GCNII: 解決了圖卷積網(wǎng)絡(luò)的過度平滑的問題，并從理論上證明了該算法的有效性。在國際知名機(jī)器學(xué)習(xí)數(shù)據(jù)網(wǎng)站PaperWithCode的各測評榜單中長期占據(jù)第一名的位置，在Cora和Pubmed數(shù)據(jù)集上的預(yù)測準(zhǔn)確率自2020年提出至今始終保持第一，獲世界人工智能大會青年優(yōu)秀論文提名獎。
分子圖基座模型Uni-Mol: 基于分子三維結(jié)構(gòu)的通用分子表征學(xué)習(xí)框架, 性能優(yōu)越、模型泛化能力強(qiáng)，在小分子性質(zhì)預(yù)測、蛋白靶點(diǎn)預(yù)測和蛋白-配體復(fù)合物構(gòu)象預(yù)測等任務(wù)上都超越之前方法。
Jittor Geometric的各種模型均以高度統(tǒng)一的模式編寫，極大降低了用戶的學(xué)習(xí)成本。在熟練掌握一類數(shù)據(jù)集和模型后，用戶可以借助Jittor Geometric自由探索各領(lǐng)域圖神經(jīng)網(wǎng)絡(luò)模型的設(shè)計，無需花費(fèi)大量時間進(jìn)行領(lǐng)域間的遷移學(xué)習(xí)。
據(jù)研發(fā)團(tuán)隊(duì)透露，未來Jittor Geometric還將進(jìn)一步拓展支持的數(shù)據(jù)集和模型，關(guān)注圖神經(jīng)網(wǎng)絡(luò)前沿研究，如大語言模型與圖學(xué)習(xí)等，助力研究人員和開發(fā)者們更便捷地進(jìn)行前沿探索和落地應(yīng)用。
多個數(shù)據(jù)集上，性能超過現(xiàn)有主流框架
Jittor Geometric開發(fā)團(tuán)隊(duì)提供了實(shí)驗(yàn)數(shù)據(jù)。以GCN模型為例，在各種經(jīng)典圖數(shù)據(jù)集上Jittor Geometric均展現(xiàn)出與現(xiàn)有主流圖機(jī)器學(xué)習(xí)框架相當(dāng)或更優(yōu)的性能。具體來說，Jittor Geometric在小規(guī)模圖數(shù)據(jù)集上的訓(xùn)練速度相較主流框架提升25%；在大規(guī)模圖數(shù)據(jù)集上顯著優(yōu)于Pytorch Geometric和Paddle Graph Learning, 與Deep Graph Library相當(dāng)。
圖 4：Jittor Geometric與各主流框架性能對比
為了方便廣大用戶上手Jittor Geometric，研發(fā)團(tuán)隊(duì)采用了與PyTorch Geometric較為相似的模塊化接口。從上手難度來看，熟悉 PyTorch Geometric的用戶可以直接上手，甚至沒有因?yàn)樽兞棵煌鴰碛绊憽?br/>
Jittor Geometric提供了良好的課程支持，幫助對圖機(jī)器學(xué)習(xí)感興趣的同學(xué)們更好地學(xué)習(xí)和設(shè)計各種圖神經(jīng)網(wǎng)絡(luò)。一方面，Jittor Geometric 定義了全面且精簡的基礎(chǔ)圖算子，同學(xué)們可以基于這些基礎(chǔ)算子實(shí)現(xiàn)多樣的消息傳遞和圖算法；另一方面還提供了各領(lǐng)域的模型實(shí)現(xiàn)、訓(xùn)練框架及數(shù)據(jù)集，幫助同學(xué)們更快上手圖機(jī)器學(xué)習(xí)領(lǐng)域的任務(wù)、數(shù)據(jù)及模型設(shè)計。
圖 5：Jittor Geometric在圖機(jī)器學(xué)習(xí)課程中的應(yīng)用
據(jù)研發(fā)團(tuán)隊(duì)介紹，Jittor Geometric已經(jīng)應(yīng)用到課程教學(xué)中，同學(xué)們基于該框架設(shè)計并實(shí)現(xiàn)了各種動態(tài)圖神經(jīng)網(wǎng)絡(luò)和異配圖神經(jīng)網(wǎng)絡(luò)，在多個數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)并作相關(guān)分析。同學(xué)們反饋該框架簡明易懂，學(xué)習(xí)成本低，在多種前沿模型及數(shù)據(jù)集上相較Torch Geometric實(shí)現(xiàn)更簡單，訓(xùn)練更快速。
Jittor Geometric的新進(jìn)展
Jittor Geometric作為一個新興圖機(jī)器學(xué)習(xí)平臺，在一些功能上，仍舊需要持續(xù)迭代完善。比如生態(tài)的建設(shè)，以及更大范圍的推廣，仍舊需要很多的努力。據(jù)研發(fā)團(tuán)隊(duì)透露，Jittor Geometric將在以下幾個方面進(jìn)行進(jìn)一步提升：
1. 更多前沿模型補(bǔ)充 拓展和完善圖神經(jīng)網(wǎng)絡(luò)領(lǐng)域的最新算法和模型，支持更多的前沿圖神經(jīng)網(wǎng)絡(luò)架構(gòu)。2. 高效動態(tài)圖計算支持 加強(qiáng)對動態(tài)圖的高效計算支持，提升圖神經(jīng)網(wǎng)絡(luò)在處理動態(tài)、時序圖數(shù)據(jù)時的性能。3. 分布式訓(xùn)練加速 進(jìn)一步優(yōu)化分布式訓(xùn)練框架，支持多節(jié)點(diǎn)、多GPU的并行計算，提升訓(xùn)練大規(guī)模圖模型的速度和效率。4. 輔助轉(zhuǎn)換腳本 提供PyTorch Geometric實(shí)現(xiàn)到Jittor Geometric實(shí)現(xiàn)的自動轉(zhuǎn)換腳本，幫助熟悉PyTorch的用戶更快上手。
研究團(tuán)隊(duì): 中國高校圖學(xué)習(xí)開源框架的新力量
據(jù)官網(wǎng)介紹，Jittor Geometric由中國人民大學(xué)魏哲巍教授和東北大學(xué)張巖峰教授及其團(tuán)隊(duì)共同打造。開發(fā)Jittor Geometric的主力，是實(shí)驗(yàn)室雷潤林，陸施展，付振波等一批博士生。他們的主要研究方向涵蓋譜域、動態(tài)、分子等多領(lǐng)域圖神經(jīng)網(wǎng)絡(luò)，圖存儲與圖計算優(yōu)化等，在重要國際會議NeurIPS, ICML, ICLR, KDD, VLDB, SIGMOD, ICDE等發(fā)表多篇論文。據(jù)雷潤林透露，團(tuán)隊(duì)接下來的重點(diǎn)將圍繞學(xué)術(shù)界展開，致力于使Jittor Geometric成為國內(nèi)外學(xué)術(shù)界最受歡迎、使用最廣泛的圖機(jī)器學(xué)習(xí)庫之一，并對AI產(chǎn)業(yè)界產(chǎn)生積極影響。
雷潤林表示：「我們聚焦于圖機(jī)器學(xué)習(xí)的前沿領(lǐng)域，關(guān)注圖學(xué)習(xí)未來發(fā)展中關(guān)鍵的動態(tài)圖和大規(guī)模圖數(shù)據(jù)處理等。同時希望能夠通過構(gòu)建活躍的開源社區(qū)，不僅完善框架本身，還形成各種圖學(xué)習(xí)的實(shí)踐指引，技術(shù)博客及問題討論等，幫助對圖機(jī)器學(xué)習(xí)領(lǐng)域感興趣的用戶更好地探索實(shí)踐。目前，Jittor Geometric已經(jīng)應(yīng)用到課堂教學(xué)中并取得良好反饋，未來還將聯(lián)合更多高校將其融入教學(xué)體系中�！�
總之，打造一款全面高效的圖機(jī)器學(xué)習(xí)庫并非易事，而選擇開源則意味著要接受來自全球開發(fā)者社區(qū)的嚴(yán)格檢驗(yàn)。
如今，Jittor Geometric已經(jīng)邁出了至關(guān)重要的一步，誠邀所有感興趣的讀者參與其中，一同見證這一新興框架的成長與發(fā)展。
參考資料：項(xiàng)目地址：https://github.com/AlgRUC/JittorGeometric
