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

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

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

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

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