Ⅰ 深度学习推理引擎的一些思考
深度学习推理引擎,特别是走编译优化路径的引擎,是解决现代深度学习应用的关键技术之一。作为曾经的传统编译器开发者,我见证了从传统编译器到深度学习领域的转变。在阿里巴巴工作期间,我积极推动基于编译优化的深度学习推理引擎的开发,这一领域在短短几年内已经取得了显着进展。
在18年初加入阿里巴巴时,基于编译优化的深度学习推理引擎并不占据主流地位,当时流行的解决方案是TFLite和NCNN等框架,它们依赖高效GEMM加速库或针对特定CPU的手写汇编代码。然而,我坚信编译器优化才是解决这个问题的最佳途径和未来方向。我的坚定信念得到了老板的支持,我们顺利推动了这个项目,最终取得了很好的落地效果。现今,TVM、GLOW、MLIR等框架已成为业界认可的编译优化深度学习推理引擎。
我选择TVM作为基础框架,原因在于它在编译优化领域与我的想法最为契合。尽管TVM在18年初并不成熟,但相比其他框架,它更加接近我的技术理想。虽然没有从零开始开发,但我认为基于已有的成熟框架进行扩展,可以更快地解决问题,同时也节省了大量时间和资源。
做深度学习推理引擎,并非易事。一方面,它涉及矩阵计算的优化,这需要对编译器、体系结构、高性能计算等领域有深刻理解。另一方面,它并非想象中那么简单,深度学习推理引擎的真正实现需要下足功夫,并涵盖量化压缩、图优化、子图分离、异构执行等多个方面。例如,在支持Hexagon DSP时,深入了解其体系结构与编程模型是关键。
对于深度学习推理引擎,接收模型是第一步,包括TensorFlow、TFLite、MXNet等。目前,这一环节在不同框架中都存在不足,模型支持程度参差不齐。以TVM为例,虽然已实现了部分算子支持,但仍有大量算子未覆盖。对于开发者而言,可选择在框架基础上进行扩展,共同推动技术进步。
在解析模型后的计算图阶段,各框架实施的优化策略不尽相同,但普遍会涉及算子融合等通用优化技术。TVM引入了类似LLVM的Pass机制,支持异构执行,但在自动化的异构子图分离方面仍有改进空间。业界对于此问题的研究,包括基于传统机器学习的Cost Model与强化学习的方法,为解决异构执行问题提供了新的思路。
计算图对接外部推理引擎是另一个值得关注的点。在实际应用中,业务部门通常追求模型快速且高效执行。以GPU为例,TensorRT提供高性能推理能力,但可能不支持某些算子。结合TVM与TensorRT,通过子图分离实现不同算子在不同硬件上的高效执行,是解决这一问题的有效途径。这一策略同样适用于NPU的支持,通过异构执行达到性能优化。
在算子层面,卷积操作是优化的重点。编译器、体系结构和高性能计算领域的专家在这方面能发挥关键作用。TVM的Auto TVM和Halide思想的应用,使得优化技术的移植和性能提升变得更为高效。编译生成部分,TVM采用LLVM为基础,简化了代码指令生成的工作量,同时提供了Tensorize机制,允许开发者自定义优化。
量化压缩是深度学习领域的重要技术之一,TVM支持不同的量化策略,包括TFLite转换和TensorRT量化。量化策略的选择和优化是关键,业界的一些框架提供快速重训练等工具,以提高模型准确度和执行效率。这方面的学习和实践对于深度学习编译器的发展具有重要意义。
总之,基于编译优化的深度学习推理引擎是未来趋势,而TVM等框架在这一领域展现出巨大潜力。随着技术的不断发展,我们可以期待更多创新和突破,为深度学习应用提供更高效、更灵活的解决方案。