科研之友
微信
新浪微博
Facebook
分享链接
摘要
随着深度学习应用的快速发展,神经网络模型的参数量变得越来越大,这意味着训练一个可用的神经网络模型需要更多的算力和更长的计算时间,因此如何提升神经网络训练的效率至关重要.然而训练效率在很大程度上取决于硬件后端和编译器.为了提升神经网络训练的性能,编译器的效率亟待提升,而这主要取决于计算图的优化、算子级别的优化和代码生成.主流的神经网络训练框架(如TensorFlow, PyTorch)使用了供应商特定的、通过手工设计算子获得的算子库.然而,手工设计算子浪费了大量的算子级别的优化空间,因此研究人员提出了TVM.作为一个端到端的编译器, TVM实现了算子级的自动优化,比现有框架进一步提高了性能.此外, TVM支持从多种神经网络框架中导入神经网络模型,并在不同主流硬件后端上部署.然而, TVM的注意力集中于提升神经网络推理任务的性能,并不支持神经网络的训练.本文提出了TVMT,第一个基于TVM、支持神经网络训练的端到端编译器.为了支持神经网络训练,本文提出了以下方法.(1)合并损失函数到现有的计算图中,以支持前向和后向传播;(2)在训练过程中采用设备到主机的机制来更新权重参数;(3)集成了最先进的张量程序调优器,以自动优化神经网络训练程序.实验结果表明,与PyTorch相比, TVMT在Intel CPU和NVIDIA GPU上的神经网络训练性能达到了最高4.88倍的提升;与TensorFlow相比, TVMT在Intel CPU和NVIDIA GPU上的神经网络训练性能达到了最高11.5倍的提升.
- 单位
