机器学习推理性能优化技术概览

ChuanXin · 发表于 2022-1-11 14:28

概述

机器学习领域最重要的两个问题是机器学习模型的生产（模型开发，模型训练，AutoML）和机器学习模型的部署（推理）。机器学习模型的生产过程关注的是模型的效果，主要由算法工程师通过模型结构设计，超参调整等方式来达成。而机器学模型的部署则关注的是模型计算的性能，主要由系统工程师通过计算图层面的优化，算子层面的优化等方式，在保证模型效果的前提之下，提升模型计算的性能。

目前，业界也产生了很多针对模型推理性能进行优化的方案：
加速库：cudnn/cublas/dnnl/zendnn等
推理框架：TensorRT/TNN/MNN/TVM等
专用硬件：各种NPU及软件栈

在诸多推理性能优化解决方案中，TensorRT由于其在GPU硬件上出色的性能优势，逐渐成为主流的推理优化方案。本文档主要对Tensorflow模型使用TensorRT进行推理性能优化的技术路径进行调研，兼顾不同方案扩展支持不同深度学习框架和推理优化后端的可复用性。

对于一个具有包容性的推理优化框架来说，需要支持不同的模型前端格式，和不同的异构计算硬件。同时，也需要支持不同的优化技术作为后端优化器。而对使用方来说，推理优化框架的输出模型结果。要么是一个统一的格式，要么仍然是原始输入格式。

统一的输出格式对应一个统一的运行时，技术上比较简洁，可以统一进行优化，可能存在开放自定义算子的成本。而原始的输入格式无需担心算子不支持问题（优化引擎不支持的算子回退给原始框架运行时即可），但是需要框架运行时本身支持优化引擎相关算子（如TF-TRT），或者提供给用户一个推理SDK(如PAI-Blade)。

而这两种技术路线，业界均有相关的方案。TF-TRT走的是保持原始输入格式的技术路，为了支持经过TF-TRT优化后模型的部署，Tensorflow框架本身集成了加载TensorRT运行时的算子。而TVM则采用的是统一格式的技术路线，经过TVM优化后的模型统一使用TVM运行时部署，而TVM运行时中集成了加载推理优化引擎（如TensorRT/MKL-DNN）的算子。

TensorRT

TensorRT原是Nvidia推出的针对GPU硬件的推理加速框架，后来也支持在CPU硬件上进行部署，其优越的性能使得TensorRT成为在服务端广泛使用的推理加速框架。TensorRT工作流程如下图所示，首先导入训练好的网络模型，然后使用优化引擎对网络模型进行算子融合，权重量化等优化操作，产出优化好的模型，最后使用TensorRT runtime执行优化好的模型。

导入模型

TensorRT支持Tensorflow，Pytorch，Caffe等多种主流的深度学习模型格式。在TensorRT早期的版本中，部署不同的深度学习框架产出的模型首先需要将模型转换为不同的中间格式，如Tensorflow模型需要转换为UFF格式，Pytorch需要转换为ONNX格式。随着后来ONNX的逐渐完善和商业上的考量，ONNX格式成为TensorRT对接不同深度学习框架的统一的模型格式。所有主流的深度学习模型在使用TensorRT优化前，统一转换为ONNX模型，然后TensorRT导入ONNX模型转换为TensorRT格式的模型，最后进行模型优化产出优化后的模型供推理服务进行部署。

使用TensorRT优化Tensorflow模型的过程如下图所示：首先，将TensorFlow计算图转换为ONNX计算图，然后将ONNX计算图转换为TRT计算图并进行优化，最后用TensorRT作为计算引擎进行推理。

这种方式可以将整个Tensorflow计算图完全转换为TRT计算图，可以做到计算图的全局优化，性能表现非常突出。

Tensorflow模型转换为ONNX模型目前已有开源工具tensorflow-onnx，而Pytorch, MxNet等框架本身提供了将模型保存为ONNX模型的接口。而ONNX模型转换为TensorRT模型也有开源工具使用onnx-tensorrt。利用这些工具，TensorRT实现了对主流训练框架模型格式的支持。

主要特性

权重与激活精度校准
支持FP32/FP16/INT8等数据格式。通过将模型量化为 INT8 来更大限度地提高吞吐量，同时保持高准确度
层与张量融合
通过融合内核中的节点，优化 GPU 显存和带宽的使用
内核自动调整
基于目标 GPU 平台选择最佳数据层和算法
动态shape
支持运行时根据张量的形状构建计算引擎
动态张量显存
更大限度减少显存占用，并高效地为张量重复利用内存
多流执行
用于并行处理多个输入流的可扩展设计

软件栈开源协议

tensorflow-onnx: Apache-2.0 License
onnx-tensorrt: Apache-2.0 License
TensorRT: Apache-2.0 License

扩展性

前端训练框架

这种方案，将onnx作为统一的计算图格式，可以比较容易地支持各种训练框架。Pytorch和MxNet框架本身集成了将模型保存为onnx格式的接口。

后端优化引擎

就后端优化引擎的扩展性而言，如果推理优化框架需要集成其他推理优化引擎如TNN，MNN等，只需要支持onnx模型格式即可。

存在问题

维护成本

深度学习模型计算图到TensorRT计算图的工作流中存在两次模型格式的转化（origin model2onnx/onnx2trt）。因此，整个转化链路中，出现不支持算子的概率较高。另外，当模型中存在用户的自定义算子时，也需要扩展开发相应的ONNX算子和TensorRT算子，否则计算图就不能转化成功。在我们接触到的客户中，自定义算子的情况并不少见。当对接的业务量增加时，需要大量人力帮助客户开发定制算子，因此业务支持效率很难随着业务量横向扩展。所以，站在算法开发者的角度来讲，出于性能考虑可以采用这种方案，但是站在需要对接大量业务的云厂商来讲，这种方案后续的维护成本较高。

优化灵活性

这种方案是将训练模型的整个计算图完全转换为推理引擎的计算图。因此，整个计算图只能使用一种推理优化框架，无法综合多种推理优化引擎的优势。

参考资料

https://docs.nvidia.com/deeplearning/tensorrt/index.html

TF-TRT

如前文所述，深度学习模型转换到TensorRT模型存在两次转换过程，存在算子不支持的可能性，需要有额外的开发支持。而TF-TRT项目正是为了解决这个问题。TF-TRT通过检测Tensorflow计算图中能够被TensorRT支持的算子构成的子图，将该子图转换为单个经过优化的TRTEngineOp节点。通过这种方式，TF-TRT既可以通过TensorRT计算引擎优化计算图推理性能，同时TensorRT不支持的算子又可以通过Tensorflow进行计算。

TF-TRT 既可以利用 TensorFlow 的灵活性，同时还可以利用可应用于 TensorRT 支持的子计算图的优化。TensorRT 只优化和执行计算图的一部分，剩余的计算图由 TensorFlow 执行。在下图所示推理示例中，TensorFlow 执行了 Reshape 运算和 Cast 运算。然后，TensorFlow 将预构建的 TensorRT 引擎 TRTEngineOp_0 的执行传递至 TensorRT 运行时。

工作流

下图展示的是使用TF-TRT优化前后的工作流对比。使用TF-TRT优化模型时，需要转化模型的计算子图，并且可选地预构建TRT引擎，然后保存优化后的计算图。最后便可以使用支持TF-TRT的Tensorflow运行时加载模型文件进行推理。

主要特性

TF-TRT可以使用TensorRT本身所具有的功能特性，如混合精度，动态Shape，层与算子融合等。另外TF-TRT还会为每个转换为TRTEngineOp的子图创建一个FunctionDef，在运行时如果执行TensorRT引擎失败，则可以回退到Tensorflow运行时。

软件栈开源协议

TF-TRT: 已经成为Tensorflow Codebase的一部分，Apache-2.0 License
TensorRT: Apache-2.0 License

扩展性

前端训练框架

TF-TRT方案本身是TensorRT针对Tensorflow框架的集成方案。如果需要支持其他训练框架模型如Pytorch/MxNet，则首先需要将TF-TRT从Tensorflow项目中解耦出来，然后开发支持解析其他训练框架模型的功能模块（或者直接支持ONNX模型格式）。
后端优化引擎

需要设计一个根据计算图特点选择不同优化引擎对不同子图进行优化的引擎。

存在问题

推理部署

TF-TRT是针对特定框架Tensorflow集成的个性化的方案，而且Tensorflow社区已经默认集成了TensorRT运行时的支持。所以经过TF-TRT优化产出的模型可以直接使用Tensorflow运行时进行推理计算。而目前社区对于Pytorch, MxNet等框架并没有类似的集成方案，Pytorch, MxNet等训练框架的运行时并没有集成TensorRT运行时的支持。因此在提供给用户优化后模型的同时，需要同时提供推理部署的SDK（类似PAI-Blade）。否则，用户无法部署优化后的推理模型。

参考资料

https://github.com/tensorflow/tensorflow/pull/16253
https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html
https://discuss.tf.wiki/t/topic/1391
https://blog.tensorflow.org/2019/06/high-performance-inference-with-TensorRT.html

TVM

TVM借鉴Halide算法和调度分离的思想，采用编译优化的技术路线对深度模型的推理性能进行优化。尤其是AutoTVM/AutoScheduler引入机器学习算法自动优化算子性能，取得了不错的效果。

BYOC

为了让算法工程师在开发新模型时不必担心性能问题，硬件后端供应商（例如 Intel、NVIDIA、ARM 等）要么提供内核库，例如 cuBLAS 或 cuDNN，要么提供带有计算图引擎的框架，如DNNL 或 TensorRT，让用户以某种方式描述他们的模型以实现高性能。此外，新兴的深度学习加速器也有自己的编译器、内核库或运行时框架，如Habana goya NPU。

但是，当用户尝试在新的内核库或设备上工作时，他们必须学习新的编程接口。因此，为了让所有用户和硬件后端提供商站在同一页面上，对统一编程接口的需求变得越来越重要。

BYOC是TVM提供的一种统一的为不同硬件后端添加kernel library(cuBLAS or cuDNN
)/compiler(deep learning accelerators)/framework(DNNL or TensorRT)的接口。不同的硬件后端或者加速库对不同的算子的加速效果是不一样的，有些算子适合用TensorRT进行性能优化，有些算子适合用DNNL进行性能优化，有些算子适合用TVM编译优化。而BYOC则可以集成不同的硬件后端，为深度学习模型计算图中的不同子图选择不同的优化组件，从而使得整个计算图的性能达到最优。

和TF-TRT方案类似，BYOC针对某个加速库筛选计算图中支持的算子构成的子图，将该子图使用该加速库进行优化，而加速库不支持的算子则回退给TVM进行编译优化。

BYOC工作流

对于一个需要优化的计算图，BYOC首先需要标注能够被指定加速库支持的算子，然后对能够支持的算子进行计算图级别的优化，如算子融合。接下来，将融合后的多个算子进一步融合为子图，以减少数据访存次数或者加速引擎加载次数。最后，为融合的子图创建一个带有Compiler属性的Relay Function。Compiler属性表明这个函数会通过指定的加速库进行性能优化。

BYOC的后端加速软件栈

目前BYOC已支持TensorRT/cuDNN/cuBlas/DNNL/Arm Compute Library等10余种针对不同硬件后端的加速库或者加速框架，支持根据模型的结构特点，灵活选择不同的加速方案优化模型的推理性能。

扩展性

TVM支持TensorRT的技术路线可以认为集合了前文两种方案的优势：
1）TVM提供了统一的Relay IR（相当于ONNX），通过Relay IR，TVM可以支持Tensorflow, Pytorch, MxNet等主流的深度学习模型。
2）基于BYOC，TVM可以集成TensorRT等不同硬件后端的软件加速栈。并且TensorRT等加速方案支持的算子可以回退给TVM进行支持。
3）BYOC提供了统一的硬件加速栈扩展接口，可以很方便地添加其他推理加速方案。

当然，如果深度学习模型中存在自定义算子或者TVM Relay IR不支持的算子，模型也无法成功转换，此时需要扩展开发TVM的Relay IR。但是对比TensorRT方案需要添加ONNX和TensorRT两层算子，TVM添加新算子的成本要低一些。因此，TVM对不同的前端和后端均具有很好的扩展性。

软件栈开源协议

TVM: Apache-2.0 License
TensorRT: Apache-2.0 License

存在问题

1. 当模型中存在自定义算子或者Relay IR不支持的算子时，需要开发支持相关算子。另外TVM的TVMDSOOp功能支持将TVM优化后的算子或者子图以custom op的形式插入到Tensorflow计算图中。在一定程度上可以缓解这个问题，但是目前TVM还不支持将整个计算图优化后保存为原始计算图格式，使用上不方便。不过可以基于TVMDSOOp进一步开发扩展。

参考资料

https://tvm.apache.org/2020/07/15/how-to-bring-your-own-codegen-to-tvm
https://tvm.apache.org/docs/deploy/tensorrt.html
https://discuss.tvm.apache.org/t/add-the-document-for-tvmdsoop/6622

		自动登录	找回密码
密码			立即注册

机器学习推理性能优化技术概览

本帖子中包含更多资源