出身背景

加速深層神經網絡（ 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=