深度学习推理引擎的一些思考

深度学习推理引擎的一些思考

我想熟悉我背景的人应该知道我以前是做传统编译器的,而加入阿里巴巴以后我开始做深度学习推理引擎,而在我们部门走编译优化这条路来做深度学习推理引擎也是我主推起来的。在很早以前的一篇文章:蓝色:手把手带你遨游TVM 我曾提到了我为什么认为这是一个编译器的问题,也为什么认为编译优化才是解决这个问题的最佳方案与未来,这也是我说服老板的技术缘由。

在18年初我刚加入阿里巴巴的时候,基于编译优化的深度学习推理引擎其实并不是主流,那时候更流行的解决方案是TFLite这样的框架,即调用高效的GEMM加速库,抑或着NCNN这样的框架,针对ARM CPU手写汇编。但是,我很坚定的对我的老板表示,这都不是未来,用编译器来解决这个问题才是未来。当时我有多坚定呢,基本上可以理解为如果不这么搞,我就离职那种。很庆幸的是,我的老板非常力挺我的想法,于是从我最开始调研到最后实际推动起来,并没有遇到很多阻力,而最后也回报了老板的信任,我们取得了很好的落地效果。而在这两年时间,无论是TVM,GLOW,MLIR等,走编译优化来做深度学习推理引擎越来越受到认可,以今年 TVM Conf Talk来看,可以看到很多公司的身影,Amazon,Facebook,Microsoft,ARM,Intel,Qualcomm,Xilinx等,当然也包括了没有列举在这上面的华为,阿里巴巴,以及很多基于TVM来做NPU或者应用的公司(题外话,据我所知,华为或许是中国投入做TVM力量最大的团队了),这里面其实很多人都是像我一样是传统编译器出身,比如参加过TVM Conf 2018的Qualcomm的Krzysztof Parzyszek,也是做LLVM编译器的,我和他就TVM如何支持Hexagon DSP也做过很多交流。说到这里,或许会问为什么我们会选取TVM作为Base。在18年初的时候,TVM其实很不成熟(而即使放到现在,我认为TVM也还有很多改善的地方,但是相比18年好很多,后文我会说到我的一些看法),但是相比那时候调研的Tensor Comprehensions, XLA, ngraph等等,TVM是最接近我想法的框架。而为什么没有完全从0开始?我当时的看法是除了开发者的技术情怀以外,我找不到任何理由从0开始,因为你最后要做出来的也是类似TVM这种,没有任何的必要,不如在这个基础上来做,很多脏活累活也都已经帮你做完了。如做编译器,现在从0开始做编译器的已经基本上少很多了,基于LLVM的基本上算是通用做法了。在底层软件这一块儿,我个人的看法是拥抱开源生态远比一个人玩强很多,一起把这个蛋糕做大。

上面说了这么多,其实还没有说到有关技术一些东西。接下来,我会就我观测的,以及实际开发的经验来说一下我的一些感想,或许会有不对的地方,权当抛砖引玉。

做深度学习推理引擎难不难?我个人认为没有想象的那么难,也没有想象的那么简单。看起来好像很废话,但是是针对不同的情景。如我刚接触的时候,我其实还是把它想象的还是很神秘的,模型跑起来后到底是怎么就把猫识别出来的呢?我现在看法就不一样了,你其实可以不管它怎么识别出来的,对于你引擎开发人员来说,这里面就是矩阵的计算,你如何把它算的又快又准就可以了,至于它为什么可以识别出来猫,可以去补补算法的知识。所以,我认为做深度学习推理引擎很适合做编译器、体系结构、HPC的人。而没有想象的那么简单是指什么呢?深度学习推理引擎要真正做好,还是需要下很多功夫的,也需要很多方面知识的。以TVM支持Hexagon DSP为例,这里面如果你不了解Hexagon的体系结构知识与编程模型,你基本上无从下手。而为了生成Hexagon DSP的指令,你还需要LLVM的知识来进行LLVM CodeGen,这又回到了传统编译器领域。而深度学习推理引擎也不仅仅如此,还包括量化压缩,图优化,子图分离,异构执行等,后续我也会谈到。

那我们从头开始说起。对于深度学习推理引擎,无论是你基于编译优化这条路,还是传统的调用加速库这条路,第一步不可避免的就是接收模型,如Tensorflow, TFLite, MXNet等。无论是TVM,还是其它框架,这方面做的其实都不够好。很多时候,框架使用者第一步尝试的就是拿你这个框架跑一下我的模型,然后发现各种算子不支持,更别提后面的性能了。如果你是基于TVM这种开源的引擎来做还好,遇到不支持的可以添加,但是如果你是使用类似高通的SNPE,Intel的OpenVINO这样的闭源方案,那你基本上除了改模型没有其他的办法。那么,对于没有开发引擎的人来说,或许很难理解这件事情,为什么各大推理引擎不能都支持完了。这里我们先限制在CNN领域的算子,其实并不是说这个多难,我个人认为这其实是一个体力活。以我做的TVM TFLite前端为例,我把这个前端做完了,并加上了convolution, relu等算子,我不能添加类似add算子的支持吗?显然不是我不能,而是我个人精力有限,我以Mobilenet V1模型为驱动,验证了我这条路和框架是通的,然后我开源出去。如果你遇到了不支持的算子,你可以在我的框架上添加,我可以帮忙Review,然后我们一起把这个壮大,这也是开源的伟大之处。然而,抛出开发者这个角度来说,这其实是最影响用户体验的一环,目前大家的做法基本上都是以用户报哪个算子不支持或者跑哪个模型发现不支持然后来增加,并不是像传统编译器一样,拿着一份表挨着挨着来做,比如Clang这样clang.llvm.org/cxx_stat

在这里就不得不提ONNX,这是很伟大的构想,如果真的成功了,我们深度学习推理引擎只需要支持ONNX即可,其余的任何框架模型都先变为ONNX即可。但是ONNX属于一把好牌打得稀烂,目前在这方面投入还算比较积极的就属于微软了,至于原因是为什么,也是很显然能想到的,这里就不再赘述了。

对于深度学习推理引擎来说,解析完模型以后都会变为自己的计算图表示,在这里面大家都会做事情,但是做的都不尽相同。不过一些通用的图优化大家都会做,如算子融合。在图这方面,TVM从NNVM变为了Relay,并引入了类似LLVM的Pass机制,并支持了异构。这方面的优势暂且不谈,因为其实你能做的,我也能做,比如异构,算子融合等。这里我谈谈我个人觉得TVM这里还可以改进的一块,那就是自动化的异构子图分离。对于设备上有多个硬件的话(以高通的硬件为例,即有CPU,GPU,DSP),如何高效的自动把计算图分离成多个计算子图,并且异构的执行,从而更高效的执行模型,这一块儿其实TVM基本上没有探索。在这方面,Training当然是有更强烈的需求,但是不是说Inference就没有。那么如果进行子图分离,哪些算子放在DSP,哪些放在GPU,哪些放在CPU,实际运行的时候,如CPU有巨大的波动如何处理等。业界有一些在基于传统机器学习的Cost Model来做这个事情,也有一些在基于强化学习来做,我觉得都是可以学习的一点。

在计算图这一块,其实还有一块儿我觉得是可以考虑的,那就是如何在计算图对接外部推理引擎。这一块儿或许听起来很矛盾,因为类似TVM这样的本就是推理引擎,你为什么要考虑对接外部推理引擎。然而我认为只有真正的去接触了业务,你才能知道大家想要的到底是什么。对于业务部门来说,想要的就两点:我的模型可以很轻易地跑起来;我的模型可以跑得很快。以GPU为例,TensorRT是一个很强的高性能推理框架,现在我业务部门有一个模型我需要你在GPU跑起来,并且我需要在XX ms 跑完。而很遗憾的是TensorRT不支持这个模型的一些算子,并且在一些层比TVM慢,一些层比TVM快。而要在规定时间完成这个业务目标,最直接的办法就是TVM + TensorRT,前端统一走TVM,然后在Relay这里进行子图分离,把一些层交给TensorRT,一些层走TVM。这样,算子都由你来控制,同时TensorRT执行它高效的一部分,TVM执行不支持的算子以及TVM本身高效的一部分从而达到目的。这一块儿其实可以推广到NPU的支持,如果NPU不暴露指令集,但是暴露了计算图的API,你不能走编译生成这条路,你也可以这样做,这样你不再受限于NPU不支持算子的问题,你可以异构的跑在TVM CPU端。

经历了计算图,我们来到算子层面。这一块儿也是大家费工夫优化的点,以CNN为例,基本上就是卷积算子。我为什么说做编译器、体系结构、HPC的同学很适合呢,就是因为我们很大一部分优化都在这里,说白了,就是如何在目标平台更高效的执行多重循环。你以前学到的Cache、Loop优化完全能无缝的派上用场。在这一块儿,我觉得各大框架都差不多,Winograd, Spatial Pack, im2col,layout变换等,大家都会做,区别点就在于谁做的更好一些,但是都半斤八两,没道理你会的我不会。针对这个,我认为TVM跟很多传统框架有几个很不同的点。第一个就是Auto TVM,这一点在业务模型中非常厉害。因为无论业务模型的卷积形状如何变,Auto TVM总能给你找到很好的参数让你的计算有很好的局部性,这一点其实很影响性能。这一点其实也被很多框架看到了,也有很多框架开始学习Auto TVM来做Auto Tuning。第二个是TVM借助Halide思想,进行了计算与调度分离,并提供了DSL。你在ARM CPU做的优化,很容易移植到x86 CPU,你不需要生写NEON与SSE,你所关注的点在于是不是应该.vectorize,是不是应该.unroll,是不是应该.parallel。这个带来的工作效率的提升,我觉得只有真正做过优化的才能体会到,人就应该干人擅长的事情,关注用哪些优化技术,而不是痛苦的把NEON移植成SSE,或者CUDA,OpenCL,Metal等。

到后面代码指令生成部分,TVM走的就是编译生成,如CPU就是走的LLVM。其实相比传统编译器,这部分的编译生成简单很多,你要去做生成的TVM IR也很少很少,如Load, Store, Add等,就几十个。如你有一个新的硬件,你从TVM IR到硬件的CodeGen工作量不会很大,如果这个硬件是基于LLVM Toolchain的,这个工作量更少,你可以复用TVM CodeGenLLVM的很大一部分。很多人会说这样的编译效率不高,我个人是持反对意见的,在这部分只要能生成你想要的指令,如mla,那么无论是走LLVM生成,还是你手写汇编生成,抑或着你去调用高效的加速库,其实没有差别。而如果你觉得做得不好,其实在TVM中提供了一个机制,叫做Tensorize,该机制允许你替代自动生成的部分,而是自己手写的代码。这是我觉得TVM在这里做得很好的地方,主推编译优化,却允许容纳手工微内核。

接下来谈谈量化压缩,对于量化有两种,一种是以TensorFlow为代表的,quantization-aware training,其完整的路线是Tensorflow -> TFLite Convert -> TFLite,其在工业界应用很广泛,还有一种是类似TensonRT这种,深度学习推理引擎做量化。我认为这两种都应该支持,如MNN目前就是支持这两种,而TVM也是。第一种抛开不谈,这种支持没有什么好说的,就是解析TFLite量化模型,然后底层需要支持INT8的运算,以及需要更高效的支持INT8的卷积运算,包括类似ARM v8.2的dot指令产生等,大家都会做。而这里想谈谈第二点,这里面其实也是很自由发挥的点,第一是如何支持类似INT N 的量化,而不仅仅是INT 8,如FPGA,其实可以跑类似INT4这样的类型。第二点是如何更好地量化,得到更好地准确度。这里面其实各家做的不同,目前TVM社区就有很多用户发帖说TVM的这一点其实做的不好,经常说准确度不好。而目前业界一些框架的做法是会加入快速的重训练,不仅仅是单纯的Post Training Quantization,这一套会抽象出来一个工具,供算法业务团队使用。而这套工具往往还会包括模型压缩,甚至告诉你怎么压缩可以在这个硬件跑的更快,而不仅仅是降低模型GFLOPS这一简单的指标,我觉得这或许都是类似TVM等引擎学习的一点。

好了,其实想说的还有很多,但是时间比较晚了,希望下次有机会再来谈谈。但是,无论如何,基于编译优化的深度学习推理引擎我认为才是最佳解决方案以及未来,而基于这一思想的解决方案终将会出现一个类似LLVM的东西出来,目前现状依然是群雄逐鹿,但是我个人觉得TVM目前是看起来最有希望的。:-)

延伸文章讨论:
谈谈对深度学习编译技术的一些思考

也谈TVM和深度学习编译器

编辑于 2019-10-24

文章被以下专栏收录