中国科大研发Capsule框架,GNN训练性能大幅提升12倍!

B站影视 2025-02-24 15:07 2

摘要:近日,中国科学技术大学苏州高等研究院的一项研究成果在国际学术界引起了广泛关注。该研究院下属的医学影像智能与机器人研究中心DDL实验室,成功研发出一种名为Capsule的新型图神经网络(GNN)训练框架,这一创新成果已被ACM International Con

近日,中国科学技术大学苏州高等研究院的一项研究成果在国际学术界引起了广泛关注。该研究院下属的医学影像智能与机器人研究中心DDL实验室,成功研发出一种名为Capsule的新型图神经网络(GNN)训练框架,这一创新成果已被ACM International Conference on Management of Data(SIGMOD)2025正式收录。

图神经网络作为一种强大的工具,在推荐系统、自然语言处理、计算化学及生物信息学等多个领域展现出了卓越的性能。现有的主流GNN训练框架,如DGL和PyG,主要依赖GPU的并行计算能力,从复杂的图数据中高效提取结构信息。然而,随着图数据规模的日益增大,GPU显存容量的限制成为制约GNN系统扩展性的关键因素。

DDL实验室针对这一难题,提出了Capsule这一核外(Out-of-Core)GNN训练框架,为大规模图数据的处理提供了全新的解决方案。Capsule通过创新的图划分和图裁剪策略,确保训练子图的结构及特征能够完全加载到GPU显存中,从而避免了反向传播过程中CPU与GPU之间的频繁I/O操作,显著提升了训练效率。

不仅如此,Capsule还采用了基于最短哈密顿回路的子图加载方式,并引入了流水线并行策略,进一步优化了系统的整体性能。这一框架的设计极具灵活性,能够无缝集成到现有的主流开源GNN训练框架中,为用户提供了极大的便利。

在实际应用中,Capsule展现出了惊人的性能提升。在针对大规模真实图数据集的测试中,Capsule在仅使用22.24%内存的情况下,相比当前最优的系统性能提升了最高可达12.02倍。Capsule还首次提供了关于训练所得嵌入方差的理论上界,为GNN的训练和应用提供了更为坚实的理论基础。

DDL实验室的这一研究成果,标志着我国在图计算系统领域取得了重大突破。Capsule框架的推出,将为社交网络分析、生物医学图谱构建等需要处理超大规模图数据的应用场景提供强有力的技术支持,推动我国相关领域的研究和应用迈向新的高度。

来源:ITBear科技资讯

相关推荐