亚马逊和华盛顿大学今天合作发布了开源的端到端深度学习编译器NNVM compiler。
先提醒一句,NNVM compiler ≠ NNVM。
NNVM是华盛顿大学博士陈天奇等人2016年发布的模块化深度学习系统,今年8月中旬,他们又推出了将深度学习工作负载部署到硬件的端到端IR堆栈TVM,也就是把深度学习模型更简单地放到各种硬件上。
当时,陈天奇把TVM+NNVM描述为“深度学习到各种硬件的完整优化工具链”,而这次推出的NNVM compiler,是一个基于TVM工具链的编译器。
项目作者之一陈天奇在微博上这样介绍这个编译器:
我们今天发布了基于TVM工具链的深度学习编译器NNVM compiler。支持将包括mxnet,pytorch,caffe2, coreml等在内的深度学习模型编译部署到硬件上并提供多级别联合优化。速度更快,部署更加轻量级。 支持包括树莓派,服务器和各种移动式设备和cuda, opencl, metal, javascript以及其它各种后端。 欢迎对于深度学习, 编译原理,高性能计算,硬件加速有兴趣的同学一起加入dmlc推动领导开源项目社区。
NNVM compiler对CoreML的支持,让开发者可以在非iOS设备上部署CoreML模型。
AI开发界的挑战
AWS AI首席科学家李沐(MXNet作者)在亚马逊博客撰文介绍称,推出这个编译器,是为了应对深度学习框架多样化为AI开发界带来的三个挑战:
一、
对于算法的开发者来说,由于各AI框架的前端交互和后端实现之间都存在很多区别,换框架很麻烦,而开发和交付过程中可能会用到的框架不止一个。
比如说有的亚马逊AWS云服务用户,为了获得EC2上的加速性能,会想要把Caffe模型部署到MXNet上。
为了应对这个问题,之前Facebook和微软也联合发布了模型间转换工具ONNX。
二、
框架的开发者需要维护多个后端,来保证自己的框架能适用于从手机芯片到数据中心GPU的各种硬件。
比如说MXNet,要支持英伟达GPU的cuDNN,还要支持英特尔CPU的MKLML。
三、
从芯片供应商的角度来看,他们每新开发一款芯片都需要支持多个AI框架,每个框架表示和执行工作负载的方式都不一样,所以,就连卷积这样一个运算,都需要用不同的方式来定义。
支持多个框架,就代表要完成巨大的工作量。
通过将框架中的深度学习模型直接部署到硬件,NNVM compiler自然也就解决了这些问题。
结构
NNVM compiler可以将前端框架中的工作负载直接编译到硬件后端,能在高层图中间表示(IR)中表示和优化普通的深度学习工作负载,也能为不同的硬件后端转换计算图、最小化内存占用、优化数据分布、融合计算模式。
编译器的典型工作流如下图所示:
这个编译器基于此前发布的TVM堆栈中的两个组件:NNVM用于计算图,TVM用于张量运算。
其中,NNVM的目标是将不同框架的工作负载表示为标准化计算图,然后将这些高级图转换为执行图。
TVM提供了一种独立于硬件的特定域语言,以简化张量索引层次中的运算符实现。另外,TVM还支持多线程、平铺、缓存等。
对框架和硬件的支持
编译器中的NNVM模块,支持下图所示的深度学习框架:
具体来说,MXNet的计算图能直接转换成NNVM图,对Keras计算图的直接支持也正在开发中。
同时,NNVM compiler还支持其他模型格式,比如说微软和Facebook前不久推出的ONNX,以及苹果CoreML。
通过支持ONNX,NNVM compiler支持Caffe2、PyTorch和CNTK框架;通过支持CoreML,这个编译器支持Caffe和Keras。
而编译器中的TVM模块,目前附带多个编码生成器,支持多种后端硬件,其中包括为X86和ARM架构的CPU生成LLVM IR,为各种GPU输出CUDA、OpenCL和Metal kernel。
性能
NNVM compiler联合使用图级和张量级优化以获得最佳性能。常规的深度学习框架会将图优化与部署runtime进行打包,而NNVM编译器将优化与实际部署运行时分离。
采用这种方法,编译的模块只需要依赖于最小的TVM runtime,当部署在Raspberry Pi或移动设备上时,只占用大约300KB。
陈天奇团队对NNVM compiler的性能进行了基准测试,并与MXNet进行了比较。这个测试基于两种典型的硬件配置:树莓派上的ARM CPU和AWS上的Nvidia GPU。
Nvidia GPU
GPU的基准和时间表由Leyuan Wang(AWS/UCDavis)和Yuwei Hu(图森)提供。他们在Nvidia K80上对NNVM编译器和MXNet进行了比较,以CUDA8和CuDNN7作为后端。这是一个非常强的基线,因为MXNet开启了从CuDNN中选择最佳内核的自动调整功能。另外,他们还使用了MXNet中优化深度内核来优化MobileNet工作负载。
如图所见,NNVM编译器生成的代码在K80上优于MXNet。这些改进源于图和内核级别的优化。值得注意的是,NNVM编译器自己升恒所有的优化GPU内核,而不需要依赖诸如CuDNN这样的外部库。
树莓派3b
树莓派编译堆栈由Ziheng Jiang(AWS/FDU)提供。他们使用OpenBLAS和NNPack对NNVM和MXNet进行了比较,尝试不同的设置来获得MXNet的最佳表现,例如为3×3卷积在NNPack中开启Winograd卷积,启动多线程,并禁用了额外的调度程序(所有的线程都被NNPack使用)。
结果如上图所示,由NNVM编译器生成的代码在ResNet18上速度快两倍。MobileNet上的差距,主要是因为现有CPU DNN库中缺乏深度卷积。NNVM编译器受益于直接生成高效的ARM代码。
开发团队
NNVM编译器的GitHub地址:
https://github.com/dmlc/nnvm
开发这个项目的依然是TVM堆栈团队,包括华盛顿大学艾伦计算机学院的陈天奇、Thierry Moreau、Haichen Shen、Luis Ceze、Carlos Guestrin和Arvind Krishnamurthy,以及亚马逊AWS AI团队的Ziheng Jiang。
另外,在TVM博客最后还鸣谢了一群社区贡献者:
在这里特别感谢Yuwen Hu(图森)、Leyuan Wang(AWS/UCDavis)、Joshua Z. Zhang(AWS)以及Xingjian Shi(HKUST)的早期贡献。我们也要感谢所有的TVM堆栈贡献者。
— 完 —