作者: 王乐园 王贻达 AWS AI
本文谈的是在移动端加速卷积神经网络。虽然AWS是个云服务公司,我们同样重视edge上的计算。现代终端设备一般都跟云端服务器相连,但只要可能,我们都希望计算可以在本地终端解决,这样做的好处是多方面的:既可以减小网络带宽的压力,又可以避免网络传输产生的时延,还可以让用户的数据更安全。现代终端设备一般用一个片上系统 (SoC)做计算,上面部署了通用的CPU和集成显卡。对于日益增多的卷积神经网络推理计算来说,在移动端的CPU(多数ARM,少数x86)上虽然优化实现相对简单(参见我们对CPU的优化),但此处它并非最佳选择,因为:1)移动端CPU算力一般弱于集成显卡(相差在2-6倍之间);2)更重要的是,已经有很多程序运行在CPU上,如果将模型推理也放在上面会导致CPU耗能过大或者CPU节流,造成耗电过快同时性能不稳定。所以在移动端进行模型计算,集成显卡是更好的选择。
说起来很有道理,但用起来就不一样了。实际中我们发现移动设备上的集成显卡利用率很低,大家并不怎么用它来跑卷积神经网络推理。原因其实很简单:难用。在AWS,我们面对很多移动端机器,里面用到集成显卡多数来自Intel, ARM和Nvidia,编程模型一般是OpenCL和CUDA。虽然对于某些特定模型和算子,硬件厂商提供了高性能库(Intel的OpenVINO, ARM的ACL, Nvidia的CuDNN),但它们覆盖度有限,用起来不灵活,造成即使对单一硬件做多模型的优化,工程代价也很大,遑论我们面对的硬件类种类繁多。总之,要用传统方法在集成显卡上实现一个通用高效的模型推理并不容易。
好在现在我们有了深度学习编译器。Apache TVM从一开始的设计目标就是做好模型和硬件的中间件,把不同模型编译到不同硬件上执行,中间尽量统一管理,这跟我们对集成显卡的需求是完全一致的。于是,我们的解法是扩展Apache TVM对集成显卡的支持,用同一套IR来统一描述神经网络并根据硬件类型lower到不同的硬件平台上去。这样我们可以尽量复用优化,并且利用TVM已有的优化方案(比如此文提到的算子优化、图优化和协同优化)来实现在集成显卡上的高效模型推理。最后,我们还在TVM上对几个在集成显卡上经常被用到的计算机视觉特定算子进行了精调。我们在三款主流厂商的集成显卡上进行了验证,包括AWS DeepLens (Intel Graphics), Acer aiSage/Rockchip RK-3399 (ARM Mali)和Nvidia Jetson Nano (Nvidia GPU with Maxell architecture)。 这个项目的结果发表在了并行计算会议ICPP 2019:A Unified Optimization Approach for CNN Model Inference on Integrated GPUs。相关代码都已经merge到Apache TVM里了。时至今日这个项目也还没有做完,我们依然在根据实际需要不断扩展对不同模型和硬件的支持。
具体地说,为了让TVM高效快速支持集成显卡上卷机神经网络的推断计算,我们在TVM已有框架下加入了下图彩色部分所示的部分,包括IR,对算子和计算图的优化,以及对CUDA/OpenCL后端的优化。
![cf0bdfafd3d9f0a74aef5b095c1b0bdf.png](https://i-blog.csdnimg.cn/blog_migrate/ec9355113fb822de3fb143d07b33eab0.jpeg)
以下是我们在不同移动端显卡上优化后的性能和原厂提供软件的比较(latency, 数字越低越好)。从图中看出我们的性能在大部分模型上比原厂快,而且模型覆盖率更高,比如在我们做实验的时候Intel OpenVINO还不支持在DeepLens上跑目标检测网络。
![1ee8c780c5d2a186840cd9af5f01aed0.png](https://i-blog.csdnimg.cn/blog_migrate/50141f6ebfd3fea8baf97a5c3c8576d0.png)
![84cf20ab2ac99e9b531a2f7ec5cfe4ab.png](https://i-blog.csdnimg.cn/blog_migrate/14cddab6e1affcf3fbb4286a66c7ab79.png)
![d7c97fa1afc0e26da7887184f052cc61.png](https://i-blog.csdnimg.cn/blog_migrate/ec70c59932492ef149867917d69f56cb.png)
下图是我们使用算子搜索和计算图搜索前后的性能变化(speedup, 数字越高越好),具体优化算法概念上跟我们在CPU上的优化是一样的,不过算子搜索模板因为硬件架构的不同而有所改变。
![125ea76242bd5656ecb2e2f2a3c4844b.png](https://i-blog.csdnimg.cn/blog_migrate/044ff712c50cec846399977f5e11fb88.png)
![86688cbefdbc53f8592e3028213b14fe.png](https://i-blog.csdnimg.cn/blog_migrate/58e1361bba19307a2c643c87f9fa3970.png)
![907714a2dc65ea7dd58475972978a1a7.png](https://i-blog.csdnimg.cn/blog_migrate/9a463d3abff855dd8e8605cc70bf7e53.png)
另外,我们也在TVM中加入了对视觉模型某些操作算子的一些细致优化。其中包括对视觉模型中常见NMS算子,扫描算子的优化。NMS中包含一个GPU上经典的优化问题:排序。排序和扫描在英伟达显卡上都已经有对应的优化,但是在英特尔集成显卡上并没有相应实现。我们用TVM IR实现了GPU经典优化,下方左图是我们对GPU排序的优化,右图是对扫描的优化。和普通的排序不同,NMS中的排序是区块排序,意思是数据分散在不同区块中,排序仅发生在每个区块中,不同区块间数据顺序无需改变。因为GPU有高速并行处理大数据的能力,不适合小区块的处理,我们将不同区块放在同一个数组里,并且记录下每个区块开始的位置。我们将这个大数组均分成n块(具体数量根据GPU硬件资源决定)。图中假设我们GPU有5个并行线程,我们将数据均分为5块。每个线程单独排序每块数据,跨越区块边界的不去排序。在每块数据排序结束后再进行合并操作,用的也是GPU经典的线程合作式合并方法。详细方法介绍请参考上方链接的论文。我们发现用TVM IR实现和用OpenCL/CUDA实现不仅无需在不同后端采用不同语言,代码数也有显著减少。
![11fec4894ef01aa0c45c7d1d87f32c8d.png](https://i-blog.csdnimg.cn/blog_migrate/e5f55ddb044bf9928243a346d3fedb9f.jpeg)
下图是计算机视觉模型在优化相应算子前后的性能比较(speedup, 数字越高越好)。
![42514a4006e9334e2b6dc3411546466c.png](https://i-blog.csdnimg.cn/blog_migrate/295fee03b114fc9120e648c0617381c2.png)
![a41caf59dc94ee357252a6a4d72a9de4.png](https://i-blog.csdnimg.cn/blog_migrate/3d115d402372efda175149461e9993a0.png)
![aa0129aee4a9002e43871e0f62df0353.png](https://i-blog.csdnimg.cn/blog_migrate/ef0ba92b7e6a55583940df180261426b.png)
虽然我们做了一系列针对集成显卡可编程性和性能的优化,还是要承认集成显卡的可编程性和灵活度比CPU要差,比如排序这种算法,在CPU上可以快速实现,但是在集成显卡上实现一个相对性能较优的方法要更大的工程量。如果用户提出一个新模型,而我们没有足够时间来实现一个集成显卡上面的优化方案,这时我们提供了一个替代方法,把新模型中暂不支持的运算放回到CPU上执行。这样可以保证任意模型我们都可以快速用我们的系统实现。实际操作中,对于一些计算量较小但控制逻辑复杂的视觉算子,在CPU上实现简单得多,而且在CPU上和集成显卡上性能差别并不是很大。例如,把SSD模型的NMS算子退回到CPU上执行与全模型在集成显卡上执行比较,在AWS Deeplens上用退回方法速度仅下降0.5%。退回的实现方法逻辑上很简单,利用了两次遍历图的逻辑。第一次遍历将需要在GPU/CPU上跑的算子进行分别标记,第二次遍历将插入数据拷贝算子。具体的退回实现方法请参看上面我们提到的论文。