今天小編分享的科技經驗:中國科大突破圖神經網絡訓練瓶頸,Capsule 系統性能提升 12 倍,歡迎閲讀。
IT 之家 2 月 23 日消息,2 月 11 日,中國科大蘇州高等研究院醫學影像智能與機器人研究中心 DDL 實驗室宣布,團隊在大規模圖神經網絡(GNN)訓練系統方面的研究論文被國際知名學術會議 ACM International Conference on Management of Data(SIGMOD)2025 錄用。
圖神經網絡(GNN)的主流訓練框架,如 DGL 和 PyG,利用 GPU 的并行處理能力從圖數據中提取結構信息,在推薦系統、自然語言處理、計算化學和生物信息學等領網域展現出卓越性能。
盡管 GPU 在 GNN 訓練中提供了強大的計算優勢,但其有限的顯存容量難以容納龐大的圖數據,因而現有 GNN 系統在大規模圖數據上的擴展性方面仍面臨挑戰。
對此,DDL 實驗室團隊提出了一種新的核外(Out-of-Core)GNN 訓練框架 —— Capsule,為大規模圖神經網絡訓練提供了高效的解決方案。
與現有的核外 GNN 框架不同,該系統通過圖劃分和圖裁剪策略将訓練子圖結構及其特征完全放入 GPU 顯存之中從而消除了反向傳播過程中 CPU 與 GPU 之間的 I/O 開銷,進而顯著提升了系統性能。
此外,Capsule 通過設計基于最短哈密頓回路的子圖加載方式和流水線并行策略,進一步優化了性能。同時,Capsule 具備即插即用的特性,能夠無縫集成至主流開源 GNN 訓練框架中。在大規模真實圖數據集上,Capsule 與現有最好的系統相比能夠在僅使用 22.24% 的内存下帶來最高 12.02 倍的性能提升,并提供了關于訓練所得嵌入方差的理論上界。
這一成果标志着我國在圖計算系統領網域取得重要突破,為社交網絡分析、生物醫學圖譜構建等需處理超大規模圖數據的場景提供了全新解決方案。
IT 之家附論文鏈接:
https://doi.org/10.1145/3709669