NVIDIA 技术博客：使用 NVIDIA QAT 工具包为 TensorFlow 和 NVIDIA TensorRT 加速量化网络

nvdev 2023-06-01 17:51:35

这篇文章最初发表在 NVIDIA 技术博客上。有关此类的更多内容，请参阅最新的 计算机视觉/视频分析 新闻和教程。

我们很高兴宣布 NVIDIA Quantization Aware Training （ QAT ） Toolkit for TensorFlow 2 目标是在 NVIDIA GPU 上使用 NVIDIA TensorRT 加速量化网络。该工具包为您提供了一个易于使用的 API ，以一种优化的方式量化网络，只需几行额外的代码即可进行 TensorRT 推理。

这篇文章伴随着走向 INT8 推理：使用 TensorRT 部署量化感知训练网络的端到端工作流 GTC 课程。有关 PyTorch 量化工具包等效工具，请参阅 PyTorch 量化 .

出身背景

加速深层神经网络（ 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 的量化工作流

有关量化、量化方法（ PTQ 与 QAT 相比）和 TensorRT 中量化的更多信息，请参阅使用 NVIDIA TensorRT 的量化感知训练实现 INT8 推理的 FP32 精度。

用于 TensorFlow 的 NVIDIA QAT 工具包

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

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

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

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

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

Feature	TFMOT	NVIDIA QAT Toolkit
QDQ node placements	Outputs and weights	Inputs and weights
Quantization support	Whole model (full) and of some layers (partial by layer class)	Extends TF quantization support: partial quantization by layer name and pattern-based layer quantization by extending CustomQDQInsertionCase
Quantization op used	Asymmetric quantization (tf.quantization.fake_quant_with_min_max_vars)	Symmetric quantization, needed for TensorRT compatibility (tf.quantization.quantize_and_dequantize_v2)

表 1 ：。 NVIDIA QAT 工具包与 TensorFlow 模型优化工具包的区别

图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”)