使用NVIDIA QAT工具包實(shí)現(xiàn)TensorRT量化網(wǎng)絡(luò)的設(shè)計(jì)

出身背景

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

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

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

量化感知訓(xùn)練

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

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

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

用于 TensorFlow 的 NVIDIA QAT 工具包

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

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

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

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

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

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

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

TensorRT 中部署 QAT 模型的工作流

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

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

假設(shè)預(yù)訓(xùn)練的 TensorFlow 2 模型為 SavedModel 格式，也稱為基線模型。

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

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

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

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

ResNet-50v1 示例

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

要求

要跟進(jìn)，您需要以下資源：

Python 3.8

TensorFlow 2.8

NVIDIA TF-QAT 工具包

TensorRT 8.4

準(zhǔn)備數(shù)據(jù)

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

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

工作流支持 TFRecord 格式，因此請(qǐng)使用以下說明（從 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)在應(yīng)該可以在$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ā)動(dòng)機(jī)

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

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

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

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