PyTorch 量化簡介

在開發機器學習應用時，有效利用伺服器端和裝置上的計算資源非常重要。為了支援在伺服器和邊緣裝置上更高效的部署，PyTorch 添加了對使用熟悉的 eager 模式 Python API 進行模型量化的支援。

量化利用 8 位整數 (int8) 指令來減小模型大小並加快推理速度（降低延遲），這對於模型能否達到服務質量目標，甚至能否適應移動裝置上的可用資源來說，都可能至關重要。即使資源不是那麼受限，它也可能使您能夠部署更大、更準確的模型。從 PyTorch 1.3 版本開始支援量化，隨著 PyTorch 1.4 的釋出，我們在 PyTorch torchvision 0.5 庫中釋出了 ResNet、ResNext、MobileNetV2、GoogleNet、InceptionV3 和 ShuffleNetV2 的量化模型。

這篇博文概述了 PyTorch 上的量化支援及其與 TorchVision 領域庫的整合。

什麼是量化？

量化是指使用較低精度資料（通常是 int8，而非浮點實現）進行計算和記憶體訪問的技術。這可以在幾個重要領域實現效能提升

• 模型大小減少 4 倍；
• 記憶體頻寬減少 2-4 倍；
• 由於節省了記憶體頻寬和 int8 算術計算速度更快，推理速度提高 2-4 倍（具體加速取決於硬體、執行時和模型）。

然而，量化並非沒有額外的成本。從根本上說，量化意味著引入近似值，導致網路的準確性略有下降。這些技術試圖最大限度地縮小全浮點精度與量化精度之間的差距。

我們將量化設計為適合 PyTorch 框架。這意味著

1. PyTorch 具有與 量化張量對應的資料型別，它們共享張量的許多功能。
2. 可以使用量化張量編寫核心，就像浮點張量的核心一樣，以自定義其實現。PyTorch 支援 `torch.nn.quantized` 和 `torch.nn.quantized.dynamic` 名稱空間中常見操作的量化模組。
3. 量化與 PyTorch 的其他部分相容：量化模型是可跟蹤和可指令碼化的。量化方法對於伺服器和移動後端幾乎相同。可以在模型中輕鬆混合量化和浮點運算。
4. 浮點張量到量化張量的對映可以透過使用者定義的觀察器/偽量化塊進行自定義。PyTorch 提供了適用於大多數用例的預設實現。

我們在 `torch.quantization` 名稱空間中開發了三種用於 PyTorch 神經網路量化的技術，作為量化工具的一部分。

PyTorch 1.3 版本開始支援的三種量化模式

1. 動態量化PyTorch 支援的最簡單的量化方法稱為動態量化。這不僅涉及將權重轉換為 int8（所有量化變體中都會發生），還涉及在計算之前（因此稱為“動態”）即時將啟用轉換為 int8。因此，計算將使用高效的 int8 矩陣乘法和卷積實現來執行，從而加快計算速度。然而，啟用以浮點格式讀取和寫入記憶體。
   • PyTorch API：我們在 PyTorch 中有一個簡單的動態量化 API。`torch.quantization.quantize_dynamic` 接受一個模型以及其他幾個引數，並生成一個量化模型！我們的端到端教程為 BERT 模型演示了這一點；雖然教程很長，包含載入預訓練模型和與量化無關的其他概念的部分，但量化 BERT 模型的部分很簡單
   import torch.quantization quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
   • 有關該函式的文件在此處，以及我們的教程中的端到端示例在此處和在此處。
2. 訓練後靜態量化透過將網路轉換為同時使用整數算術和 int8 記憶體訪問，可以進一步提高效能（延遲）。靜態量化執行額外的步驟，即首先將批次資料饋送到網路中並計算不同啟用的最終分佈（具體來說，這是透過在不同點插入“觀察器”模組來記錄這些分佈來完成的）。此資訊用於確定在推理時應如何具體量化不同的啟用（一種簡單的方法是將啟用的整個範圍簡單地劃分為 256 個級別，但我們也支援更復雜的方法）。重要的是，此額外步驟允許我們在操作之間傳遞量化值，而不是在每次操作之間將這些值轉換為浮點數——然後再轉換回整數——從而顯著加速。透過此版本，我們支援多項功能，允許使用者最佳化其靜態量化
   1. 觀察器：您可以自定義觀察器模組，這些模組指定在量化之前如何收集統計資料，以嘗試更高階的方法來量化資料。
   2. 運算元融合：您可以將多個操作融合為一個操作，從而節省記憶體訪問，同時提高操作的數值精度。
   3. 逐通道量化：我們可以獨立量化卷積/線性層中每個輸出通道的權重，這可以在幾乎相同的速度下實現更高的精度。
   • PyTorch API:
     • 要融合模組，我們有 `torch.quantization.fuse_modules`
     • 觀察器使用 `torch.quantization.prepare` 插入
     • 最後，量化本身使用 `torch.quantization.convert` 完成
   我們有一個包含端到端量化示例的教程（這個教程也涵蓋了我們的第三種量化方法，即量化感知訓練），但由於我們簡單的 API，對預訓練模型 `myModel` 執行訓練後靜態量化的三行程式碼是：`# 為伺服器 (x86) 部署設定量化配置 myModel.qconfig = torch.quantization.get_default_config('fbgemm') # 插入觀察器 torch.quantization.prepare(myModel, inplace=True) # 校準模型並收集統計資訊 # 轉換為量化版本 torch.quantization.convert(myModel, inplace=True)`
3. 量化感知訓練量化感知訓練 (QAT) 是第三種方法，也是通常能帶來最高準確率的方法。在 QAT 中，所有權重和啟用在訓練的前向和後向傳播期間都進行“偽量化”：也就是說，浮點值被四捨五入以模擬 int8 值，但所有計算仍然使用浮點數完成。因此，訓練期間的所有權重調整都是在“感知”到模型最終將被量化的事實的情況下進行的；因此，量化後，這種方法通常比其他兩種方法產生更高的準確率。
   • PyTorch API:
     • `torch.quantization.prepare_qat` 將偽量化模組插入模型量化中。
     • 與靜態量化 API 類似，`torch.quantization.convert` 在訓練完成後實際量化模型。
   例如，在端到端示例中，我們載入了一個預訓練模型作為 `qat_model`，然後我們簡單地使用以下程式碼執行量化感知訓練：`# 為 QAT 指定量化配置 qat_model.qconfig=torch.quantization.get_default_qat_qconfig('fbgemm') # 準備 QAT torch.quantization.prepare_qat(qat_model, inplace=True) # 轉換為量化版本，移除 dropout，以檢查每個 epoch 的準確性 quantized_model=torch.quantization.convert(qat_model.eval(), inplace=False)`

裝置和運算元支援

量化支援僅限於可用運算元的一部分，具體取決於所使用的方法，有關支援的運算元列表，請參閱文件：https://pytorch.com.tw/docs/stable/quantization.html。

可用運算元集和量化數值也取決於用於執行量化模型的後端。目前，量化運算元僅在以下後端（x86 和 ARM）中支援 CPU 推理。量化配置（如何量化張量）和量化核心（量化張量的算術運算）都與後端相關。可以透過以下方式指定後端：

import torchbackend='fbgemm'
# 'fbgemm' for server, 'qnnpack' for mobile
my_model.qconfig = torch.quantization.get_default_qconfig(backend)
# prepare and convert model
# Set the backend on which the quantized kernels need to be run
torch.backends.quantized.engine=backend

然而，量化感知訓練是在全浮點模式下進行的，可以在 GPU 或 CPU 上執行。量化感知訓練通常僅用於 CNN 模型，當訓練後靜態或動態量化無法達到足夠的準確性時。這可能發生在高度最佳化以實現小尺寸的模型（如 Mobilenet）中。

在 torchvision 中的整合

我們還在 torchvision 中對一些最流行的模型啟用了量化：Googlenet、Inception、Resnet、ResNeXt、Mobilenet 和 Shufflenet。我們以三種形式將這些更改上游到 torchvision

1. 預訓練的量化權重，以便您可以立即使用它們。
2. 量化就緒的模型定義，以便您可以進行訓練後量化或量化感知訓練。
3. 一個用於進行量化感知訓練的指令碼——它適用於任何這些模型，儘管您將在下面瞭解到，我們發現只有在使用 Mobilenet 實現準確性時才需要它。
4. 我們還有一個教程，展示瞭如何使用 torchvision 模型之一進行量化遷移學習。

選擇方法

方案的選擇取決於多個因素

1. 模型/目標要求：某些模型可能對量化敏感，需要進行量化感知訓練。
2. 運算元/後端支援：某些後端需要完全量化的運算元。

目前，運算元覆蓋率有限，可能會限制下表中列出的選擇：下表提供了指導。

模型型別	首選方案	原因
LSTM/RNN	動態量化	吞吐量受權重計算/記憶體頻寬主導
BERT/Transformer	動態量化	吞吐量受權重計算/記憶體頻寬主導
CNN	靜態量化	吞吐量受啟用記憶體頻寬限制
CNN	量化感知訓練	在靜態量化無法達到準確性要求的情況下

效能結果

與浮點實現相比，量化可使模型大小減少 4 倍，速度提高 2 到 3 倍，具體取決於硬體平臺和基準測試模型。一些示例結果如下：

模型	浮點延遲 (ms)	量化延遲 (ms)	推理效能提升	裝置	注意事項
BERT	581	313	1.8x	Xeon-D2191 (1.6GHz)	批處理大小 = 1，最大序列長度 = 128，單執行緒，x86-64，動態量化
Resnet-50	214	103	2x	Xeon-D2191 (1.6GHz)	單執行緒，x86-64，靜態量化
Mobilenet-v2	97	17	5.7x	三星 S9	靜態量化，浮點數基於 Caffe2 執行時，未最佳化

精度結果

我們還將靜態量化模型與 Imagenet 上的浮點模型進行了精度比較。對於動態量化，我們比較了 BERT 在 GLUE MRPC 基準測試上的 F1 分數。

計算機視覺模型準確性

模型	Top-1 準確率（浮點）	Top-1 準確率（量化）	量化方案
Googlenet	69.8	69.7	靜態訓練後量化
Inception-v3	77.5	77.1	靜態訓練後量化
ResNet-18	69.8	69.4	靜態訓練後量化
Resnet-50	76.1	75.9	靜態訓練後量化
ResNext-101 32x8d	79.3	79	靜態訓練後量化
Mobilenet-v2	71.9	71.6	量化感知訓練
Shufflenet-v2	69.4	68.4	靜態訓練後量化

語音和自然語言處理模型準確性

模型	F1 (GLUEMRPC) 浮點	F1 (GLUEMRPC) 量化	量化方案
BERT	0.902	0.895	動態量化

結論

要開始在 PyTorch 中量化模型，請從 PyTorch 網站上的教程開始。如果您正在處理序列資料，請從 LSTM 的動態量化或 BERT 開始。如果您正在處理影像資料，我們建議從 量化遷移學習教程開始。然後您可以探索 靜態訓練後量化。如果您發現訓練後量化的準確率下降過高，請嘗試 量化感知訓練。

如果您遇到問題，可以透過在 discuss.pytorch.org 釋出問題來獲得社群幫助，請使用量化類別來解決與量化相關的問題。

此文章由 Raghuraman Krishnamoorthi、James Reed、Min Ni、Chris Gottbrath 和 Seth Weidman 撰寫。特別感謝 Jianyu Huang、Lingyi Liu 和 Haixin Liu 提供本文中包含的量化指標。

延伸閱讀:

1. PyTorch 在 Neurips 上的量化演示：(https://research.fb.com/wp-content/uploads/2019/12/2.-Quantization.pptx)
2. 量化張量 (https://github.com/pytorch/pytorch/wiki/ Introducing-Quantized-Tensor)
3. Github 上的量化 RFC (https://github.com/pytorch/pytorch/ issues/18318)