量化番外篇 | TensorRT-8的量化细节

作者 | Oldpan  编辑 | oldpan博客

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心【模型部署】技术交流群

好久不见各位~

这篇文章很久之前写完一直没有整理，最近终于是整理差不多了，赶紧发出来。

本文接着《必看部署系列-神经网络量化教程：第一讲！》这一篇接着来说。上一篇主要说了量化的一些基本知识、为啥要量化以及基本的对称量化这些概念知识点。按理说应该继续讲下非对称量化、量化方式等等一些细节，不过有一段时间在做基于TensorRT的量化，需要看下TensorRT的量化细节，就趁这次机会讲一下。

算是量化番外篇。

这是偏实践的一篇，主要过一下TensorRT对于explict quantization的流程和通用的量化思路。

0x01 TensorRT量化

都2022年了，量化技术已经很成熟了，各种量化框架^[1]和量化算法层出不穷。我之前接触过几个量化框架，大部分都是在算法层面模拟一下，实际上无法直接部署到具体的硬件层，也只是停留在算法的层面。而现在成熟的量化框架已经不少，开源的也有很多，无论是pytorch、TVM还是TensorRT，基于这些框架的GPU和CPU量化已经应用了不少，我也看了看最近商汤新开源的量化框架ppq，同样也挺成熟了，最起码用起来是的的确确可以实际部署，为我们带来性能的提升。

上一篇主要是理论细节比较多，那么这一篇主要说说实际的量化流程。要实际用起来、跑起来才有意义。因为有一段时间在用TensorRT，所以就说说TensorRT的量化细节和实际量化流程吧！

TensorRT的量化工具也比较成熟了。支持PTQ和QAT量化，官方也提供了一些工具^[2]去帮助我们实现量化（无论是基于trt本身还是基于周边工具）。

当然除了TensorRT我也用过一些其他的量化框架，也写过一些代码。其实大部分量化方式基本大同小异，大方向都是读取模型，转化为IR进行图分析，做一些优化策略等等，关于怎么组织图，怎么优化结构可能会不一样。还有具体的校准算法的不同，不过总体上，量化的整体思路是差不多的。

因此，了解TensorRT的量化过程是是挺重要的，也有助于理解其他框架的量化方式，毕竟万变不离其宗。

0x02 TensorRT的量化模式

TensorRT有两种量化模式，分别是implicitly以及explicitly量化。前者是隐式量化，在trt7版本之前用的比较多。而后者显式量化是在8版本后才完全支持，具体就是可以加载带有QDQ信息的模型然后生成对应量化版本的engine。

两种量化模型的一些支持情况：

与隐式量化相关性较强的是训练后量化。

训练后量化

训练后量化即PTQ量化，trt的训练后量化算法第一次公布在2017年，NVIDIA放出了使用交叉熵量化的一个PPT，简单说明了其量化原理和流程，其思想集成在trt内部供用户去使用。对我们是闭源的，我们只能通过trt提供的API去量化。

不需要训练，只需要提供一些样本图片，然后在已经训练好的模型上进行校准，统计出来需要的每一层的scale就可以实现量化了，大概流程就是这样：

具体使用就是，我们导出ONNX模型，转换为TensorRT的过程中可以使用trt提供的Calibration方法去校准，这个使用起来比较简单。可以直接使用trt官方提供的trtexec命令去实现，也可以使用trt提供的python或者C++的API接口去量化，比较容易。

目前，TensorRT提供的后训练量化算法也多了好多，分别适合于不同的任务：

• EntropyCalibratorV2

Entropy calibration chooses the tensor’s scale factor to optimize the quantized tensor’s information-theoretic content, and usually suppresses outliers in the distribution. This is the current and recommended entropy calibrator and is required for DLA. Calibration happens before Layer fusion by default. It is recommended for CNN-based networks.

• MinMaxCalibrator

This calibrator uses the entire range of the activation distribution to determine the scale factor. It seems to work better for NLP tasks. Calibration happens before Layer fusion by default. This is recommended for networks such as NVIDIA BERT (an optimized version of Google's official implementation).

• EntropyCalibrator

This is the original entropy calibrator. It is less complicated to use than the LegacyCalibrator and typically produces better results. Calibration happens after Layer fusion by default.

• LegacyCalibrator

This calibrator is for compatibility with TensorRT 2.0 EA. This calibrator requires user parameterization and is provided as a fallback option if the other calibrators yield poor results. Calibration happens after Layer fusion by default. You can customize this calibrator to implement percentile max, for example, 99.99% percentile max is observed to have best accuracy for NVIDIA BERT.

通过上述这些算法量化时，TensorRT会在优化网络的时候尝试INT8精度，假如某一层在INT8精度下速度优于默认精度（FP32或者FP16）则优先使用INT8。这个时候我们无法控制某一层的精度，因为TensorRT是以速度优化为优先的（很有可能某一层你想让它跑int8结果却是fp32）。即使我们使用API去设置也不行，比如set_precision这个函数，因为TensorRT还会做图级别的优化，它如果发现这个op（显式设置了INT8精度）和另一个op可以合并，就会忽略你设置的INT8精度。

说白了就是不好控制。我也尝试过这种方式，简单情况，简单模型问题不大（resnet系列），涉及到比较复杂的（transformer）这个设置精度可能不管用，谁知道TensorRT内部是怎么做优化的呢，毕竟是黑盒子。

训练中量化

训练中量化(QAT)是TensorRT8新出的一个“新特性”，这个特性其实是指TensorRT有直接加载QAT模型的能力。QAT模型这里是指包含QDQ操作的量化模型。实际上QAT过程和TensorRT没有太大关系，trt只是一个推理框架，实际的训练中量化操作一般都是在训练框架中去做，比如我们熟悉的Pytorch。（当然也不排除之后一些优化框架也会有训练功能，因此同样可以在优化框架中做）

TensorRT-8可以显式地load包含有QAT量化信息的ONNX模型，实现一系列优化后，可以生成INT8的engine。

QAT量化信息的ONNX模型长这样：

可以看到有QuantizeLiner和DequantizeLiner模块，也就是对应的QDQ模块，包含了该层或者该激活值的量化scale和zero-point。QDQ模块会参与训练，负责将输入的FP32张量量化为INT8，随后再进行反量化将INT8的张量在变为FP32。实际网络中训练使用的精度还是FP32，只不过这个量化算子在训练中可以学习到量化和反量化的尺度信息，这样训练的时候就可以让模型权重和量化参数更好地适应量化这个过程（当然，scale参数也是可以学习的），量化后的精度也相对更高一些。

QAT量化中最重要的就是fake量化算子，fake算子负责将输入该算子的参数和输入先量化后反量化，然后记录这个scale，就是模拟上图这个过程。

比如我们有一个网络，精度是FP32，输入和权重因此也是FP32：

我们可以插入fake算子：

FQ(fake-quan)算子会将FP32精度的输入和权重转化为INT8再转回FP32，记住转换过程中的尺度信息。

这些fake-quan算子在ONNX中可以表示为QDQ算子：

什么是QDQ呢，QDQ就是Q（量化）和DQ（反量化）两个op，在网络中通常作为模拟量化的op，比如：

输入X是FP32类型的op，输出是FP32，然后在输入A这个op时会经过Q（即量化）操作，这个时候操作A我们会默认是INT8类型的操作，A操作之后会经过DQ（即反量化）操作将A输出的INT8类型的结果转化为FP32类型的结果并传给下一个FP32类型的op。

那么QDQ有啥用呢？

• 第一个是可以存储量化信息，比如scale和zero_point啥的，这些信息可以放在Q和QD操作中

• 第二个可以当做是显式指定哪一层是量化层，我们可以默认认为包在QDQ操作中间的op都是INT8类型的op，也就是我们需要量化的op

比如下图，可以通过QDQ的位置决定每一层OP的精度：

因此对比显式量化（explicitly），trt的隐式量化（implicitly）就没有那么直接，在trt-8版本之前我们一般都是借助trt的内部的量化算法去量化，在构建engine的时候传入图像进行校准，执行的是训练后量化的过程。

而有了QDQ信息，TensorRT在解析模型的时候会根据QDQ的位置找到（我们给予提示的）可量化的op，然后与QDQ融合（吸收尺度信息到OP中）：

融合后该算子就是实打实的INT8算子，我们也可以通过调整QDQ的位置来设置网络每一个op的精度（某些op必须高精度，因此QDQ的位置要放对）：

也可以显式地插入QDQ告诉TensorRT哪些层是INT8，哪些层可以被fuse：

经过一系列融合优化后，最终生成量化版的engine：

总得来说，TensorRT加载QAT的ONNX模型并且优化的整理流程如下：

因为TensorRT8可以直接加载通过QTA量化后且导出为ONNX的模型，官方也提供了Pytorch量化配套工具，可谓是一步到位。

TensorRT的量化性能是非常好的，可能有些模型或者op已经被其他库超越（比如openppl或者tvm），不过TensorRT胜在支持的比较广泛，用户很多，大部分模型都有前人踩过坑，经验相对较多些，而且支持dynamic shape，适用的场景也较多。

不过TensorRT也有缺点，就是自定义的INT8插件不是很好搞，很多坑要踩，也就是自己添加新的支持难度稍大一些。对于某些层不支持或者有bug的情况，除了在issue中催一下官方尽快更新之外，也没有其他办法了。

各个层对INT8的支持

在官方文档的Layer specific restrictions这一节中有详细的说明，常见的卷积、反卷积、BN、矩阵乘法等等都是支持的，更多可以自己去查：

传送门：

• https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html

显式量化相关的TensorRT层

TensorRT显式量化主要参与的op是IQuantizeLayer和IDequantizeLayer这俩，即Q和DQ。在构建TensorRT-network的时候就可以通过这两个op来控制网络的量化细节。

IQuantizeLayer

这个层就是将浮点型的Tensor转换为，通过add_quantize这个API添加：

• 执行output = clamp(round(input / scale) + zeroPt)

• Clamping is in the range [-128, 127]

• API参考：https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#iquantizelayer

IDequantizeLayer

与IQuantizeLayer作用相反，通过add_dequantize添加。

• 执行声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop】