【儀表網(wǎng) 研發(fā)快訊】2月11日，中國科大蘇州高等研究院醫(yī)學影像智能與機器人研究中心DDL實驗室團隊(Data Darkness Lab)在大規(guī)模圖神經(jīng)網(wǎng)絡(luò)訓練系統(tǒng)方面的研究論文“Capsule: an Out-of-Core Training Mechanism for Colossal GNNs”被國際知名學術(shù)會議ACM International Conference on Management of Data(SIGMOD)2025錄用。
 

示意圖：Capsule系統(tǒng)工作流框架
 

　　圖神經(jīng)網(wǎng)絡(luò)(GNN)的主流訓練框架，如DGL和PyG，利用GPU的并行處理能力從圖數(shù)據(jù)中提取結(jié)構(gòu)信息，在推薦系統(tǒng)、自然語言處理、計算化學和生物信息學等領(lǐng)域展現(xiàn)出卓越性能。盡管GPU在GNN訓練中提供了強大的計算優(yōu)勢，但其有限的顯存容量難以容納龐大的圖數(shù)據(jù)，因而現(xiàn)有GNN系統(tǒng)在大規(guī)模圖數(shù)據(jù)上的擴展性方面仍面臨挑戰(zhàn)。對此，DDL實驗室團隊提出了一種新的核外(Out-of-Core)GNN訓練框架——Capsule，為大規(guī)模圖神經(jīng)網(wǎng)絡(luò)訓練提供了高效的解決方案。與現(xiàn)有的核外GNN框架不同，該系統(tǒng)通過圖劃分和圖裁剪策略將訓練子圖結(jié)構(gòu)及其特征完全放入GPU顯存之中從而消除了反向傳播過程中CPU與GPU之間的I/O開銷，進而顯著提升了系統(tǒng)性能。此外，Capsule通過設(shè)計基于最短哈密頓回路的子圖加載方式和流水線并行策略，進一步優(yōu)化了性能。同時，Capsule具備即插即用的特性，能夠無縫集成至主流開源GNN訓練框架中。在大規(guī)模真實圖數(shù)據(jù)集上，Capsule與現(xiàn)有最好的系統(tǒng)相比能夠在僅使用22.24%的內(nèi)存下帶來最高12.02倍的性能提升，并提供了關(guān)于訓練所得嵌入方差的理論上界。
 

　　計算機科學與技術(shù)學院碩士生向泳安和人工智能與數(shù)據(jù)科學學院博士生丁澤中為論文共同第一作者，通訊作者為生物醫(yī)學工程學院謝希科研究員，其他合作者還包括人工智能與數(shù)據(jù)科學學院碩士生郭睿和王上游以及生物醫(yī)學工程學院周少華教授。中國科學技術(shù)大學為唯一單位。DDL實驗室已連續(xù)兩年成功在SIGMOD上發(fā)表關(guān)于圖計算系統(tǒng)的高水平研究成果，顯示了在該領(lǐng)域的研究能力。
 

　　注：SIGMOD是數(shù)據(jù)管理領(lǐng)域國際公認的最高水平學術(shù)會議之一，由美國計算機協(xié)會(ACM)數(shù)據(jù)管理專業(yè)委員會在1975年發(fā)起并舉辦，具有悠久的歷史和很高的影響力。在中國計算機學會(CCF)推薦的“數(shù)據(jù)庫/數(shù)據(jù)挖掘/內(nèi)容檢索”A類學術(shù)會議中，SIGMOD位居首位。