利用 NVIDIA TensorRT 量化感知訓練實現 INT8 推理的 FP32 精度

深度學習正在徹底改變行業提供產品和服務的方式。這些服務包括用于計算機視覺的對象檢測、分類和分割，以及用于基于語言的應用程序的文本提取、分類和摘要。這些應用程序必須實時運行。

大多數模型都采用浮點 32 位算法進行訓練，以利用更大的動態范圍。然而，在推理時，這些模型可能需要更長的時間來預測結果相比，精度降低推理，造成一些延遲的實時響應，并影響用戶體驗。

在許多情況下，最好使用精度降低的整數或 8 位整數。挑戰在于訓練后簡單地四舍五入權重可能導致較低的模型精度，特別是當權重具有較大的動態范圍時。本文簡單介紹了量化感知訓練（ QAT ），以及如何在訓練過程中實現偽量化，并用 NVIDIA TensorRT 8 . 0 進行推理。

概述

模型量化是一種流行的深度學習優化方法，其中模型數據（包括網絡參數和激活）從浮點表示轉換為較低精度表示，通常使用 8 位整數。這有幾個好處：

在處理 8 位整數數據時， NVIDIA GPU 使用更快更便宜的 8 位張量核來計算卷積和矩陣乘法運算。這會產生更多的計算吞吐量，這在計算受限的層上尤其有效。
將數據從內存移動到計算元素（在 NVIDIA GPU s 中的流式多處理器）需要時間和精力，而且還會產生熱量。將激活和參數數據的精度從 32 位浮點值降低到 8 位整數可導致 4 倍的數據縮減，從而省電并減少產生的熱量。
有些層有帶寬限制（內存有限）。這意味著它們的實現將大部分時間用于讀寫數據，因此減少它們的計算時間并不會減少它們的總體運行時間。帶寬限制層從減少的帶寬需求中獲益最大。
減少內存占用意味著模型需要更少的存儲空間，參數更新更小，緩存利用率更高，等等。

量子化方法

量化有很多好處，但是參數和數據精度的降低很容易影響模型的任務精度。考慮到 32 位浮點可以在區間[-3 . 4e38 ， 3 . 40e38]中表示大約 40 億個數字。這個可表示數的區間也被稱為 dynamic-range 。兩個相鄰的可表示數字之間的距離是表示的精確。

浮點數在動態范圍內分布不均勻，大約一半的可表示浮點數在區間[-1 , 1]內。換言之，[-1 ， 1]區間中的可表示數字將比[1 ， 2]中的數字具有更高的精度。[-1 ， 1]中可表示的 32 位浮點數的高密度有助于深度學習模型，其中參數和數據的大部分分布質量都在零附近。

但是，使用 8 位整數表示法，只能表示 2⁸不同的價值觀。這 256 個值可以均勻地或不均勻地分布，例如，為了在零附近獲得更高的精度。所有主流的深度學習硬件和軟件都選擇使用統一的表示，因為它能夠使用高吞吐量的并行或矢量化整數數學管道進行計算。

轉換浮點張量的表示 ( $x_{f}$ ) 到 8 位表示 ( $x_{q}$ ) ，使用 scale-factor 將浮點張量的動態范圍映射到[-128 ， 127]：

$x_{q} = Clip(Round(x_{f}/scale))$

這是對稱量子化，因為動態范圍是關于原點對稱的。 $Round$ 是一個應用某些舍入策略將有理數舍入為整數的函數；和 $Clip$ 是一個函數，用于剪裁超出[-128127]區間的異常值。 TensorRT 使用對稱量化來表示激活數據和模型權重。

圖 1 的頂部是一個任意浮點張量的圖 $x_{f}$ ，表示為元素分布的直方圖。我們選擇了一個對稱的系數范圍來表示量化張量：[ $-amax$ , $amax$ ]. 在這里 $amax$ 是要表示的絕對值最大的元素。要計算量化比例，請將浮點動態范圍劃分為 256 個相等部分：

$amax = max(abs(x_{f}))$

$scale = (2 * amax) / 256$

這里顯示的計算比例的方法使用 全量程 ，可以用有符號 8 位整數表示：[-128 ， 127]。 TensorRT 顯式精度（ Q / DQ ）網絡在量化權重和激活時使用此范圍。

使用 8 位整數表示的動態范圍與舍入操作引入的誤差之間存在緊張關系。較大的動態范圍意味著原始浮點張量中的更多值用量化張量表示，但也意味著使用較低的精度和引入較大的舍入誤差。

選擇較小的動態范圍可以減小舍入誤差，但會引入剪裁誤差。超出動態范圍的浮點值將被剪裁為動態范圍的最小/最大值。

The diagram shows the range of the floating point integer on the top and its mapping to 8 bit signed integer between -128 to 127. — *圖 1 .浮點張量的 8 位有符號整數量化 $x_{f}$* *對稱動態范圍* $x_{f} [-amax, amax]$ *通過量化映射到* [-128 ， 127]。

*圖 1 .浮點張量的 8 位有符號整數量化 $x_{f}$* *對稱動態范圍* $x_{f} [-amax, amax]$ *通過量化映射到* [-128 ， 127]。

為了解決精度損失對任務精度的影響，人們發展了各種量化技術。這些技術可以分為兩類：訓練后量化（ PTQ ）或量化感知訓練（ QAT ）。

顧名思義， PTQ 是在高精度模型經過訓練后進行的。使用 PTQ ，量化權重很容易。您可以訪問權重張量并可以測量它們的分布。量化激活更具挑戰性，因為激活分布必須使用實際輸入數據進行測量。

為此，使用代表任務實際輸入數據的小數據集評估訓練的浮點模型，并收集有關層間激活分布的統計信息。作為最后一步，使用幾個優化目標之一確定模型激活張量的量化尺度。這個過程是校準，使用的代表數據集是 calibration-dataset 。

有時 PTQ 不能達到可接受的任務準確性。這是當你 MIG HT 考慮使用 QAT 的時候。 QAT 背后的思想很簡單：如果在訓練階段包含量化誤差，就可以提高量化模型的精度。它使網絡能夠適應量化的權值和激活。

有各種各樣的方法來執行 QAT ，從一個未經訓練的模型開始到一個預先訓練的模型開始。通過在訓練圖中插入假量化操作來模擬數據和參數的量化，所有的方法都改變了訓練方案，將量化誤差包含在訓練損失中。這些運算被稱為“假”運算，因為它們對數據進行量化，然后立即對數據進行去量化，這樣運算的計算就保持浮點精度。這個技巧在深度學習框架中增加了量化噪聲而沒有太大變化。

在前向過程中，對浮點權重和激活進行偽量化，并使用這些偽量化的權重和激活來執行層的操作。在向后過程中，使用權重的漸變來更新浮點權重。為了處理量化梯度，除了未定義的點之外，幾乎所有地方都是零，您可以使用（直通估計器（ STE ），它通過偽量化操作符傳遞梯度。當 QAT 過程完成時，偽量化層持有量化尺度，您可以使用這些尺度來量化模型用于推理的權重和激活。

Fake Quantization operators are added in the forward to introduce training loss which get optimized and reduced in the backward pass. — *圖 2 .訓練前傳球中的 QAT 偽量化算子* （左） *后傳球* （右）

PTQ 是這兩種方法中比較流行的方法，因為它簡單，不涉及訓練管道，這也使得它成為一種更快的方法。然而， QAT 幾乎總是產生更好的精度，有時這是唯一可以接受的方法。

TensorRT 中的量子化

TensorRT 8 . 0 支持使用兩種不同處理模式的 INT8 模型。第一種處理模式使用 TensorRT 張量動態范圍 API ，并利用 INT8 精度（ 8 位有符號整數）計算和數據機會優化推理延遲。

The workflow on the left shows the steps that Post Training Quantization requires, and on the right the workflow for Quantization Aware Training is shown which is much simpler. — *圖 3 . TensorRT PTQ 工作流程* （左） *vs . TensorRT INT8 量子化，使用從配置張量動態范圍導出的量子化尺度* （右）

當 TensorRT 執行完整的 PTQ 校準配方時，以及當 TensorRT 使用預配置的張量動態范圍時，使用此模式（圖 3 ）。另一種 TensorRT INT8 處理模式用于處理具有 QuantizeLayer/DequantizeLayer 層的浮點 ONNX 網絡，并遵循顯式量化規則。有關差異的更多信息，請參閱 TensorRT 開發人員指南中的顯式量化與 PTQ 處理。

TensorRT 量化工具箱

TensorRT ZCK4 的量化工具箱通過提供一個方便的 PyTorch 庫來補充 TensorRT ，該庫有助于生成可優化的 QAT 模型。該工具包提供了一個 API 來自動或手動為 QAT 或 PTQ 準備模型。

API 的核心是 TensorQuantizer 模塊，它可以量化、偽量化或收集張量的統計信息。它與 QuantDescriptor 一起使用，后者描述了如何量化張量。在 TensorQuantizer 之上分層的是量化模塊，這些模塊被設計為 PyTorch 全精度模塊的替代品。這些是使用 TensorQuantizer 對模塊的權重和輸入進行偽量化或收集統計信息的方便模塊。

API 支持將 PyTorch 模塊自動轉換為其量化版本。轉換也可以使用 API 手動完成，這允許在不想量化所有模塊的情況下進行部分量化。例如，一些層可能對量化更敏感，并且使其未量化可提高任務精度。

This includes tools for both PTQ as well as QAT such as Calibrator, TensorQuantizer, QuantDescriptor. — *圖 4 . TensorRT 量化工具箱組件*

在 NVIDIA 量子化白皮書中詳細描述了 QAT 的 TensorRT 特定配方，其中包括對量化方法的更嚴格的討論，以及在各種學習任務上比較 QAT 和 PTQ 的實驗結果。

代碼示例演練

本節描述了工具箱中包含的分類任務量化例子。

QAT 的推薦工具箱配方要求從預訓練模型開始，因為展示已經指出，從預訓練模型開始并進行微調可以獲得更好的精度，并且需要的迭代次數要少得多。在本例中，加載一個預訓練 ResNet50 模型。從 bash shell 運行示例的命令行參數：

python3 classification_flow.py --data-dir [path to ImageNet DS] --out-dir . --num-finetune-epochs 1 --evaluate-onnx --pretrained --calibrator=histogram --model resnet50_res

--data-dir 參數指向 ImageNet （ ILSVRC2012 ）數據集，您必須分別使用 download 數據集。 --calibrator=histogram 參數指定在微調模型之前，應該使用直方圖校準器對模型進行校準。其余的參數以及更多的參數都記錄在示例中。

ResNet50 模型最初來自 Facebook 的 Torchvision 包，但是因為它包含一些重要的更改（跳過連接的量化），所以網絡定義包含在工具箱中（ resnet50_res ）。有關詳細信息，請參閱 Q / DQ 層布置建議。

下面是代碼的簡要概述。有關詳細信息，請參閱量化 ResNet50 。

 # Prepare the pretrained model and data loaders
 model, data_loader_train, data_loader_test, data_loader_onnx = prepare_model(
 ??????? args.model_name,
 ??????? args.data_dir,
 ??????? not args.disable_pcq,
 ??????? args.batch_size_train,
 ??????? args.batch_size_test,
 ??????? args.batch_size_onnx,
 ??????? args.calibrator,
 ??????? args.pretrained,
 ??????? args.ckpt_path,
 ??????? args.ckpt_url)

函數 prepare_model 像往常一樣實例化數據加載器和模型，但它也配置量化描述符。舉個例子：

 # Initialize quantization
 if per_channel_quantization:
 ??????? quant_desc_input = QuantDescriptor(calib_method=calibrator)
 else:
 ??????? quant_desc_input = QuantDescriptor(calib_method=calibrator, axis=None)
 quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
 quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
 quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
 quant_desc_weight = QuantDescriptor(calib_method=calibrator, axis=None)
 quant_nn.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
 quant_nn.QuantConvTranspose2d.set_default_quant_desc_weight(quant_desc_weight)
 quant_nn.QuantLinear.set_default_quant_desc_weight(quant_desc_weight)

QuantDescriptor 的實例描述了如何通過配置校準方法和量化軸來校準和量化張量。對于每個量化操作（例如 quant_nn.QuantConv2d ），您可以在 QuantDescriptor 中分別配置激活和權重，因為它們使用不同的偽量化節點。

然后在訓練圖中添加假量化節點。下面的代碼（ quant_modules.initialize ）在幕后動態地修補 PyTorch 代碼，以便將 torch.nn.module 的一些子類替換為它們的量化對應項，實例化模型的模塊，然后還原動態修補程序（ quant_modules.deactivate ）。例如，將 torch.nn.conv2d 替換為 pytorch_quantization.nn.QuantConv2d ，其在執行 2D 卷積之前執行偽量化。應該在模型實例化之前調用方法 quant_modules.initialize 。

quant_modules.initialize()
model = torchvision.models.__dict__[model_name](pretrained=pretrained)
quant_modules.deactivate()

接下來，收集校準數據的統計信息（ collect_stats ）：將校準數據饋送到模型，并以直方圖的形式收集每個層的激活分布統計信息以進行量化。收集直方圖數據后，使用一個或多個校準算法（ compute_amax ）校準刻度（ calibrate_model ）。

在標定過程中，盡量確定每一層的量化尺度，以達到優化模型精度等目標。目前有兩種校準器等級：

pytorch_quantization.calib.histogram – 使用熵最小化（ KLD ）、均方誤差最小化（ MSE ）或百分位度量方法（選擇動態范圍，以表示指定的分布百分比）。
pytorch_quantization.calib.max – 使用最大激活值進行校準（表示浮點數據的整個動態范圍）。

要在以后確定校準方法的質量，請在數據集上評估模型精度。該工具包可以很容易地比較四種不同校準方法的結果，以發現適用于特定模型的最佳方法。該工具包可以擴展專有的校準算法。有關更多信息，請參閱 ResNet50 示例筆記本。

如果模型的精度令人滿意，你不必繼續進行 QAT 。您可以導出到 ONNX 并完成。這就是 PTQ 配方。 TensorRT 給出了具有量化尺度的 Q / DQ 算子的 ONNX 模型，并優化了模型進行推理。所以，這是一個 PTQ 工作流，它產生了一個 Q / dqonnx 模型。

要繼續到 QAT 階段，請選擇最佳校準、量化模型。使用 QAT 對原始訓練計劃的 10% 左右進行微調，并使用退火學習率計劃，最后導出到 ONNX 。有關更多信息，請參閱深度學習推理的整數量化：原理與實證評價白皮書。

導出到 ONNX 時，需要記住以下幾點：

ONNX opset 13 中引入了每通道量化（ PCQ ），因此如果您按照建議使用 PCQ ，請注意您使用的 opset 版本。
參數 do_constant_folding 應設置為 True ，以生成可讀性更好的較小模型。

torch.onnx.export(model, dummy_input, onnx_filename, verbose=False, opset_version=opset_version, do_constant_folding=True)

當模型最終導出到 ONNX 時，偽量化節點作為兩個獨立的 ONNX 操作符導出到 ONNX ： QuantizeLinear 和 DequantizeLinear （如圖 5 中的 Q 和 DQ 所示）。

The fake quantization nodes added in the training framework which is PyTorch in this case gets converted and exported to ONNX as two separate ONNX Operators as QuantizeLinear and DequantizeLinear. — *圖 5 .偽量化運算符被轉換為* Q/DQ *PyTorch 模型導出到 ONNX 時的 ONNX 運算符*

QAT 推斷階段

在高層次上， TensorRT 使用 Q / DQ 運算符處理 ONNX 模型，類似于 TensorRT 處理任何其他 ONNX 模型的方式：

TensorRT 導入包含 Q / DQ 操作的 ONNX 模型。
它執行一組專門用于 Q / DQ 處理的優化。
它繼續執行常規優化過程。
它為推理執行構建了一個特定于平臺的執行計劃文件。此計劃文件包含量化操作和權重。

除了啟用 INT8 外，在 TensorRT 中構建 Q / DQ 網絡不需要任何特殊的生成器配置，因為在網絡中檢測到 Q / DQ 層時，它會自動啟用。使用 TensorRT 示例應用程序 trtexec 構建 Q / DQ 網絡的最小命令如下：

$ trtexec -int8 <onnx file>

TensorRT 使用稱為 顯式量子化 的特殊模式優化 Q / DQ 網絡，這是出于對網絡處理可預測性的要求和對用于網絡操作的算術精度的控制。處理可預測性是保持原始模型計算精度的保證。其思想是， Q / DQ 層指定必須發生精度轉換的位置，并且所有優化必須保留原始 ONNX 模型的算術語義。

對比 TensorRT Q / DQ 處理和普通 TensorRT INT8 處理有助于更好地解釋這一點。在 plain TensorRT 中，使用動態范圍 API 或通過校準過程為 INT8 網絡張量分配量化尺度。 TensorRT 在應用后端優化時將模型視為浮點模型，并使用 INT8 作為另一個工具來優化層執行時間。如果一個層在 INT8 中運行得更快，那么它被配置為使用 INT8 。否則，使用 FP32 或 FP16 ，以較快者為準。在這種模式下， TensorRT 只針對延遲進行優化，您幾乎無法控制量化哪些操作。

相反，在顯式量化中， Q / DQ 層指定必須發生精度轉換的位置。優化器不允許執行非由網絡指定的精度轉換。即使這樣的轉換提高了層精度（例如，選擇 FP16 實現而不是 INT8 實現），并且即使這樣的轉換會導致執行速度更快的計劃文件（例如，在 V100 上， INT8 不被張量核加速時，首選 INT8 而不是 FP16 ），這也是正確的。

在顯式量化中，您可以完全控制精度轉換，并且量化是可預測的。 TensorRT 仍然優化性能，但要保持原始模型的算術精度。不支持在 Q / DQ 網絡上使用動態范圍 API 。

顯式量化優化過程分為三個階段：

首先，優化器嘗試最大化模型的 INT8 數據，并使用 Q / DQ 層傳播進行計算。 Q / DQ 傳播是一組規則，指定 Q / DQ 層如何在網絡中 MIG 速率。例如， QuantizeLayer 可以 MIG 通過與 ReLU 激活層交換位置來對網絡的開始部分進行速率調整。通過這樣做， ReLU 層的輸入和輸出激活減少到 INT8 精度，帶寬需求減少 4 倍。
然后，優化器融合層來創建對 INT8 輸入操作的量化操作，并使用 INT8 數學管道。例如， QuantizeLayer 可以與 ConvolutionLayer 融合。
最后， TensorRT 自動調諧器優化器搜索每一層的最快實現，同時也尊重該層指定的輸入和輸出精度。

有關 TensorRT 執行的主要顯式量化優化的更多信息，請參閱 TensorRT 開發者指南。

通過構建 TensorRT Q / DQ 網絡創建的計劃文件包含量化的權重和操作，可以部署。 EfficientNet 是需要 QAT 來保持準確性的網絡之一。下表比較了 PTQ 和 QAT 。