出身背景

加速深層神經(jīng)網(wǎng)絡（ DNN ）推理是實現(xiàn)實時應用（如圖像分類、圖像分割、自然語言處理等）延遲關鍵部署的重要步驟。

改進 DNN 推理延遲的需要引發(fā)了人們對以較低精度運行這些模型的興趣，如 FP16 和 INT8 。在 INT8 精度下運行 DNN 可以提供比其浮點對應項更快的推理速度和更低的內(nèi)存占用。 NVIDIA TensorRT 支持訓練后量化（ PTQ ）和 QAT 技術(shù)，將浮點 DNN 模型轉(zhuǎn)換為 INT8 精度。

在這篇文章中，我們討論了這些技術(shù)，介紹了用于 TensorFlow 的 NVIDIA QAT 工具包，并演示了一個端到端工作流，以設計最適合 TensorRT 部署的量化網(wǎng)絡。

量化感知訓練

QAT 背后的主要思想是通過最小化訓練期間的量化誤差來模擬低精度行為。為此，可以通過在所需層周圍添加量化和去量化（ QDQ ）節(jié)點來修改 DNN 圖。這使得量化網(wǎng)絡能夠?qū)⒂捎谀Ｐ土炕统瑓?shù)的微調(diào)而導致的 PTQ 精度損失降至最低。

另一方面， PTQ 在模型已經(jīng)訓練之后，使用校準數(shù)據(jù)集執(zhí)行模型量化。由于量化沒有反映在訓練過程中，這可能導致精度下降。圖 1 顯示了這兩個過程。

圖 1 通過 PTQ 和 QAT 的量化工作流

用于 TensorFlow 的 NVIDIA QAT 工具包

該工具包的目標是使您能夠以最適合于 TensorRT 部署的方式輕松量化網(wǎng)絡。

目前， TensorFlow 在其開源軟件 模型優(yōu)化工具包 中提供非對稱量化。他們的量化方法包括在所需層的輸出和權(quán)重（如果適用）處插入 QDQ 節(jié)點，并提供完整模型或部分層類類型的量化。這是為 TFLite 部署而優(yōu)化的，而不是 TensorRT 部署。

需要此工具包來獲得一個量化模型，該模型非常適合 TensorRT 部署。 TensorRT optimizer 傳播 Q 和 DQ 節(jié)點，并通過網(wǎng)絡上的浮點操作將它們?nèi)诤显谝黄?，以最大?INT8 中可以處理的圖形比例。這將導致 NVIDIA GPU 上的最佳模型加速。我們的量化方法包括在所需層的輸入和權(quán)重（如果適用）處插入 QDQ 節(jié)點。

我們還執(zhí)行對稱量化（ TensorRT 使用），并通過層名稱和 基于模式的層量化 的部分量化提供擴展量化支持。

表 1 總結(jié)了 TFMOT 和用于 TensorFlow 的 NVIDIA QAT 工具包之間的差異。

圖2顯示了一個簡單模型的前/后示例，用 Netron 可視化。QDQ節(jié)點放置在所需層的輸入和權(quán)重（如適用）中，即卷積（Conv）和完全連接（MatMul）。

圖 2 量化前后的模型示例（分別為基線和 QAT 模型）

TensorRT 中部署 QAT 模型的工作流

圖 3 顯示了在 TensorRT 中部署 QAT 模型的完整工作流，該模型是通過 QAT 工具包獲得的。

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

假設預訓練的 TensorFlow 2 模型為 SavedModel 格式，也稱為基線模型。

使用quantize_model功能對該模型進行量化，該功能使用 QDQ 節(jié)點克隆并包裝每個所需的層。

微調(diào)獲得的量化模型，在訓練期間模擬量化，并將其保存為SavedModel格式。

將其轉(zhuǎn)換為 ONNX 。

然后， TensorRT 使用 ONNX 圖來執(zhí)行層融合和其他圖優(yōu)化，如 專用 QDQ 優(yōu)化 ，并生成一個用于更快推理的引擎。

ResNet-50v1 示例

在本例中，我們將向您展示如何使用 TensorFlow 2 工具包量化和微調(diào) QAT 模型，以及如何在 TensorRT 中部署該量化模型。有關更多信息，請參閱完整的 example_resnet50v1.ipynb Jupyter 筆記本。

要求

要跟進，您需要以下資源：

Python 3.8

TensorFlow 2.8

NVIDIA TF-QAT 工具包

TensorRT 8.4

準備數(shù)據(jù)

對于本例，使用 ImageNet 2012 數(shù)據(jù)集 進行圖像分類（任務 1 ），由于訪問協(xié)議的條款，需要手動下載。 QAT 模型微調(diào)需要此數(shù)據(jù)集，它還用于評估基線和 QAT 模型。

登錄或注冊鏈接網(wǎng)站，下載列車/驗證數(shù)據(jù)。您應該至少有 155 GB 的可用空間。

工作流支持 TFRecord 格式，因此請使用以下說明（從 TensorFlow 說明 ） 轉(zhuǎn)換下載的。將 ImageNet 文件轉(zhuǎn)換為所需格式：

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.

您現(xiàn)在應該可以在$IMAGENET_HOME中看到兼容的數(shù)據(jù)集。

量化和微調(diào)模型

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 轉(zhuǎn)換為 ONNX

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

部署 TensorRT 發(fā)動機

將 ONNX 模型轉(zhuǎn)換為 TensorRT 引擎（還可以獲得延遲測量）：

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

獲取驗證數(shù)據(jù)集的準確性結(jié)果：

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