PyTorch 2 匯出量化用於 OpenVINO torch.compile 後端¶

作者: Daniil Lyakhov, Aamir Nazir, Alexander Suslov, Yamini Nimmagadda, Alexander Kozlov

先決條件¶

介紹¶

注意

這是一項實驗性功能，量化 API 可能會發生變化。

本教程演示瞭如何在 PyTorch 2 匯出量化流程中使用 Neural Network Compression Framework (NNCF) 中的 OpenVINOQuantizer，為 OpenVINO torch.compile 後端生成定製的量化模型，並解釋瞭如何將量化模型降低到 OpenVINO 表示形式。 OpenVINOQuantizer 透過放置專為 OpenVINO 設計的量化器，充分發揮了低精度 OpenVINO 核心的潛力。

PyTorch 2 匯出量化流程使用 torch.export 將模型捕獲到圖中，並在 ATen 圖之上執行量化轉換。此方法有望實現更高的模型覆蓋率、更強的靈活性和簡化的使用者體驗。OpenVINO 後端將 TorchDynamo 生成的 FX 圖編譯為最佳化的 OpenVINO 模型。

量化流程主要包括四個步驟

步驟 1：基於 torch 匯出機制，從即時模式（eager）模型中捕獲 FX 圖。
步驟 2：基於捕獲的 FX 圖，使用 OpenVINOQuantizer 應用 PyTorch 2 匯出量化流程。
步驟 3：使用 torch.compile API 將量化模型降低到 OpenVINO 表示形式。
可選步驟 4：透過 quantize_pt2e 方法改進量化模型的指標。

這個流程的高階架構可能如下所示

float_model(Python)                          Example Input
    \                                              /
     \                                            /
—--------------------------------------------------------
|                         export                       |
—--------------------------------------------------------
                            |
                    FX Graph in ATen
                            |
                            |           OpenVINOQuantizer
                            |                 /
—--------------------------------------------------------
|                      prepare_pt2e                     |
|                           |                           |
|                       Calibrate
|                           |                           |
|                      convert_pt2e                     |
—--------------------------------------------------------
                            |
                     Quantized Model
                            |
—--------------------------------------------------------
|                  Lower into Inductor                  |
—--------------------------------------------------------
                            |
                      OpenVINO model

訓練後量化¶

現在，我們將透過一個分步教程，向您展示如何將它與 torchvision resnet18 模型一起用於訓練後量化。

先決條件：OpenVINO 和 NNCF 安裝¶

可以透過 pip 發行版輕鬆安裝 OpenVINO 和 NNCF

pip install -U pip
pip install openvino nncf

1. 捕獲 FX 圖¶

我們將首先執行必要的匯入，從即時模式（eager）模組中捕獲 FX 圖。

import openvino.torch
import torch
import torchvision.models as models
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e

import nncf.torch

# Create the Eager Model
model_name = "resnet18"
model = models.__dict__[model_name](pretrained=True)

# Set the model to eval mode
model = model.eval()

# Create the data, using the dummy data here as an example
traced_bs = 50
x = torch.randn(traced_bs, 3, 224, 224)
example_inputs = (x,)

# Capture the FX Graph to be quantized
with torch.no_grad():
    exported_model = torch.export.export(model, example_inputs).module()

2. 應用量化¶

捕獲要量化的 FX 模組後，我們將匯入 OpenVINOQuantizer。

from nncf.experimental.torch.fx import OpenVINOQuantizer

quantizer = OpenVINOQuantizer()

OpenVINOQuantizer 有幾個可選引數，允許調整量化過程以獲得更準確的模型。以下是基本引數及其描述列表：

preset - 定義模型的量化方案。有兩種預設可用：
- PERFORMANCE (預設) - 定義權重和啟用的對稱量化
- MIXED - 權重使用對稱量化，啟用使用非對稱量化。此預設推薦用於具有非 ReLU 和非對稱啟用函式（例如 ELU、PReLU、GELU 等）的模型。
OpenVINOQuantizer(preset=nncf.QuantizationPreset.MIXED)
model_type - 用於指定特定模型型別所需的量化方案。Transformer 是唯一支援的特殊量化方案，用於在量化 Transformer 模型（BERT、Llama 等）後保留其準確性。None 為預設值，即未定義特定方案。
OpenVINOQuantizer(model_type=nncf.ModelType.Transformer)

ignored_scope - 此引數可用於將某些層排除在量化過程之外，以保留模型準確性。例如，當您想將模型的最後一層排除在量化之外時。以下是一些如何使用此引數的示例：

#Exclude by layer name:
names = ['layer_1', 'layer_2', 'layer_3']
OpenVINOQuantizer(ignored_scope=nncf.IgnoredScope(names=names))

#Exclude by layer type:
types = ['Conv2d', 'Linear']
OpenVINOQuantizer(ignored_scope=nncf.IgnoredScope(types=types))

#Exclude by regular expression:
regex = '.*layer_.*'
OpenVINOQuantizer(ignored_scope=nncf.IgnoredScope(patterns=regex))

#Exclude by subgraphs:
# In this case, all nodes along all simple paths in the graph
# from input to output nodes will be excluded from the quantization process.
subgraph = nncf.Subgraph(inputs=['layer_1', 'layer_2'], outputs=['layer_3'])
OpenVINOQuantizer(ignored_scope=nncf.IgnoredScope(subgraphs=[subgraph]))

target_device - 定義目標裝置，在最佳化過程中將考慮其特定性。支援以下值：ANY (預設)、CPU、CPU_SPR、GPU 和 NPU。
OpenVINOQuantizer(target_device=nncf.TargetDevice.CPU)

有關 OpenVINOQuantizer 的更多詳細資訊，請參閱文件。

匯入特定於後端的 Quantizer 後，我們將準備模型以進行訓練後量化。 prepare_pt2e 將 BatchNorm 運算元摺疊到前面的 Conv2d 運算元中，並在模型中的適當位置插入觀察器。

prepared_model = prepare_pt2e(exported_model, quantizer)

現在，在將觀察器插入模型後，我們將校準 prepared_model。

# We use the dummy data as an example here
prepared_model(*example_inputs)

最後，我們將校準後的模型轉換為量化模型。convert_pt2e 接受一個校準過的模型並生成一個量化模型。

quantized_model = convert_pt2e(prepared_model, fold_quantize=False)

完成這些步驟後，我們就完成了量化流程的執行，並將得到量化模型。

3. 降低到 OpenVINO 表示形式¶

之後，FX 圖可以使用 torch.compile(…, backend=”openvino”) 功能利用 OpenVINO 最佳化。

with torch.no_grad():
    optimized_model = torch.compile(quantized_model, backend="openvino")

    # Running some benchmark
    optimized_model(*example_inputs)

最佳化後的模型使用的是專為 Intel CPU 設計的低階核心。與易於理解的模型相比，這應該會顯著加快推理速度。

4. 可選：改進量化模型指標¶

NNCF 實現了一些高階量化演算法，例如 SmoothQuant 和 BiasCorrection，這些演算法有助於在最小化原始模型和壓縮模型之間輸出差異的同時，提高量化模型的指標。這些高階 NNCF 演算法可以透過 NNCF 的 quantize_pt2e API 訪問。

from nncf.experimental.torch.fx import quantize_pt2e

calibration_loader = torch.utils.data.DataLoader(...)


def transform_fn(data_item):
    images, _ = data_item
    return images


calibration_dataset = nncf.Dataset(calibration_loader, transform_fn)

with torch.no_grad():
    exported_model = torch.export.export(model, example_inputs).module()

quantized_model = quantize_pt2e(
    exported_model, quantizer, calibration_dataset, smooth_quant=True, fast_bias_correction=False
)

有關更多詳細資訊，請參閱文件以及關於 Resnet18 量化的完整示例。

結論¶

本教程介紹瞭如何將 torch.compile 與 OpenVINO 後端和 OpenVINO 量化器一起使用。有關 NNCF 和 PyTorch 模型 NNCF 量化流程的更多詳細資訊，請參閱 NNCF 量化指南。有關其他資訊，請檢視透過 torch.compile 進行 OpenVINO 部署文件。