OclDNN:一种可应用于TensorFlow的通用DNN库

作者:陈锐; 孙羽菲; 郭强; 隋轶丞; 周振辉; 石昌青; 张玉志
来源:计算机工程, 2023, 49(04): 138-148.
DOI:10.19678/j.issn.1000-3428.0064142

摘要

深度学习模型的构建、训练以及推理离不开TensorFlow等机器学习框架中深度学习算子的支撑,对于卷积、池化等深度学习中被高频调用或计算量较大的算子,机器学习框架一般通过调用深度神经网络(DNN)库来提升计算效能。现有DNN库主要由英伟达、AMD等少数国外厂商开发并根据自有硬件设备特点进行优化,但其封闭性导致其他厂商生产的通用加速器难以在深度学习领域发挥作用。为解决现有DNN库无法支持国产加速器的问题,使得深度学习模型能够调用国产加速器进行运算,研究跨平台的通用DNN库,通过对开源MIOpen的结构特点和调用方式进行分析,提出修改和重构该库的方法,并实现一种基于OpenCL的DNN(OclDNN)库。考虑到TensorFlow较高的流行度及其对DNN库调用的特殊性与复杂性,研究通用DNN库在TensorFlow中的集成方法,通过StreamExecutor中的OpenCL平台实现对OclDNN的调用。实验结果表明,OclDNN在英伟达、华为等不同厂商的计算设备上运算结果正确可靠,在相同实验环境下,深度学习算子使用OclDNN时的加速性能比传统CPU并行算法提升了5~60倍。

全文