使用NVIDIA QAT工具包实现TensorRT量化网络的设计

　　出身背景

　　加速深层神经网络（ DNN ）推理是实现实时应用（如图像分类、图像分割、自然语言处理等）延迟关键部署的重要步骤。

　　改进 DNN 推理延迟的需要引发了人们对以较低精度运行这些模型的兴趣，如 FP16 和 INT8 。在 INT8 精度下运行 DNN 可以提供比其浮点对应项更快的推理速度和更低的内存占用。 NVIDIA TensorRT 支持训练后量化（ PTQ ）和 QAT 技术，将浮点 DNN 模型转换为 INT8 精度。

　　在这篇文章中，我们讨论了这些技术，介绍了用于 TensorFlow 的 NVIDIA QAT 工具包，并演示了一个端到端工作流，以设计最适合 TensorRT 部署的量化网络。

　　量化感知训练

　　QAT 背后的主要思想是通过最小化训练期间的量化误差来模拟低精度行为。为此，可以通过在所需层周围添加量化和去量化（ QDQ ）节点来修改 DNN 图。这使得量化网络能够将由于模型量化和超参数的微调而导致的 PTQ 精度损失降至最低。

　　另一方面， PTQ 在模型已经训练之后，使用校准数据集执行模型量化。由于量化没有反映在训练过程中，这可能导致精度下降。图 1 显示了这两个过程。

图 1 通过 PTQ 和 QAT 的量化工作流

　　用于 TensorFlow 的 NVIDIA QAT 工具包

　　该工具包的目标是使您能够以最适合于 TensorRT 部署的方式轻松量化网络。

　　目前， TensorFlow 在其开源软件 模型优化工具包 中提供非对称量化。他们的量化方法包括在所需层的输出和权重（如果适用）处插入 QDQ 节点，并提供完整模型或部分层类类型的量化。这是为 TFLite 部署而优化的，而不是 TensorRT 部署。

　　需要此工具包来获得一个量化模型，该模型非常适合 TensorRT 部署。 TensorRT optimizer 传播 Q 和 DQ 节点，并通过网络上的浮点操作将它们融合在一起，以最大化 INT8 中可以处理的图形比例。这将导致 NVIDIA GPU 上的最佳模型加速。我们的量化方法包括在所需层的输入和权重（如果适用）处插入 QDQ 节点。

　　我们还执行对称量化（ TensorRT 使用），并通过层名称和 基于模式的层量化 的部分量化提供扩展量化支持。

　　表 1 总结了 TFMOT 和用于 TensorFlow 的 NVIDIA QAT 工具包之间的差异。

　　图2显示了一个简单模型的前/后示例，用 Netron 可视化。QDQ节点放置在所需层的输入和权重（如适用）中，即卷积（Conv）和完全连接（MatMul）。

　　图 2 量化前后的模型示例（分别为基线和 QAT 模型）

　　TensorRT 中部署 QAT 模型的工作流

　　图 3 显示了在 TensorRT 中部署 QAT 模型的完整工作流，该模型是通过 QAT 工具包获得的。

　　图 3 TensorRT 使用 QAT 工具包获得的 QAT 模型的部署工作流

　　假设预训练的 TensorFlow 2 模型为 SavedModel 格式，也称为基线模型。

　　使用quantize_model功能对该模型进行量化，该功能使用 QDQ 节点克隆并包装每个所需的层。

　　微调获得的量化模型，在训练期间模拟量化，并将其保存为SavedModel格式。

　　将其转换为 ONNX 。

　　然后， TensorRT 使用 ONNX 图来执行层融合和其他图优化，如 专用 QDQ 优化 ，并生成一个用于更快推理的引擎。

　　ResNet-50v1 示例

　　在本例中，我们将向您展示如何使用 TensorFlow 2 工具包量化和微调 QAT 模型，以及如何在 TensorRT 中部署该量化模型。有关更多信息，请参阅完整的 example_resnet50v1.ipynb Jupyter 笔记本。

　　要求

　　要跟进，您需要以下资源：

　　Python 3.8

　　TensorFlow 2.8

　　NVIDIA TF-QAT 工具包

　　TensorRT 8.4

　　准备数据

　　对于本例，使用 ImageNet 2012 数据集 进行图像分类（任务 1 ），由于访问协议的条款，需要手动下载。 QAT 模型微调需要此数据集，它还用于评估基线和 QAT 模型。

　　登录或注册链接网站，下载列车/验证数据。您应该至少有 155 GB 的可用空间。

　　工作流支持 TFRecord 格式，因此请使用以下说明（从 TensorFlow 说明 ） 转换下载的。将 ImageNet 文件转换为所需格式：

　　set IMAGENET_HOME=/path/to/imagenet/tar/files in data/imagenet_data_setup.sh 。

　　将 imagenet_to_gcs.py 下载到$IMAGENET_HOME。

　　Run 。/data/imagenet_data_setup.sh.

　　您现在应该可以在$IMAGENET_HOME中看到兼容的数据集。

　　量化和微调模型

from tensorflow_quantization import quantize_model from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase # Create baseline model model = tf.keras.applications.ResNet50(weights="imagenet", classifier_activation="softmax") # Quantize model q_model = quantize_model(model, custom_qdq_cases=[ResNetV1QDQCase()]) # Fine-tune q_model.compile( optimizer="sgd", loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=["accuracy"] ) q_model.fit( train_batches, validation_data=val_batches, batch_size=64, steps_per_epoch=500, epochs=2 ) # Save as TF 2 SavedModel q_model.save(“saved_model_qat”)

　　将 SavedModel 转换为 ONNX

$ python -m tf2onnx.convert --saved-model=
       
        --output=
        
         --opset 13
        
       

　　部署 TensorRT 发动机

　　将 ONNX 模型转换为 TensorRT 引擎（还可以获得延迟测量）：

$ trtexec --onnx=
       
        --int8 --saveEngine=
        
         -v
        
       

　　获取验证数据集的准确性结果：

$ python infer_engine.py --engine=
       
        --data_dir=
        
         -b=