GSO:基于图神经网络的深度学习计算图子图替换优化框架

• Published in: 计算机科学 Vol. 49; no. 3; pp. 86 - 91
• Main Authors: 苗旭鹏, 周跃, 邵蓥侠, 崔斌
• Format: Journal Article
• Language: Chinese
• Published: 北京大学信息科学技术学院 北京 100871%北京邮电大学计算机学院 北京 100871 2022
• Subjects: 计算图优化;子图替换;深度学习;图神经网络
• ISSN: 1002-137X
• DOI: 10.11896/jsjkx.210700199

TP311; 深度学习在各种实际应用中取得了巨大成功,如何有效提高各种复杂的深度学习模型在硬件设备上的执行效率是该领域重要的研究内容之一.深度学习框架通常将深度学习模型表达为由基础算子构成的计算图,为了提高计算图的执行效率,传统的深度学习系统通常基于一些专家设计的子图替换规则,采用启发式搜索算法来优化计算图.它们的不足主要有:1)搜索空间大,效率低下;2)缺乏可拓展性;3)难以利用历史优化结果.为了解决上述问题,文中提出了GSO,即一个基于图神经网络的深度学习计算图子图替换优化框架.该框架将计算图的子图优化建模成经典的子图匹配问题,基于计算图中算子的特征信息和计算图的拓扑结构信息,通过图神经网络模型来估计每种子图替换规则的匹配可行性和位置.基于与主流深度学习系统兼容的Python接口实现了GSO,实验结果表明:1)相比全量的子图替换规则,基于图神经网络的子图匹配预测可以最多减少92％的搜索空间;2)相比现有的启发式搜索算法,GSO可以更快地完成计算图子图替换优化(2倍以上),并使优化后的子图最多得到34％的加速.