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 导出机制 从动态模型捕获 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¶
OpenVINO 和 NNCF 可通过 pip 分发方式 轻松安装：
pip install -U pip
pip install openvino, nncf
1. 捕获 FX 图¶
我们将从必要的导入开始，从动态模块中捕获 FX 图。
import copy
import openvino.torch
import torch
import torchvision.models as models
from torch.ao.quantization.quantize_pt2e import convert_pt2e
from torch.ao.quantization.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(), nncf.torch.disable_patching():
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(), nncf.torch.disable_patching():
optimized_model = torch.compile(quantized_model, backend="openvino")
# Running some benchmark
optimized_model(*example_inputs)
优化后的模型使用专为英特尔 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)
quantized_model = quantize_pt2e(
exported_model, quantizer, calibration_dataset, smooth_quant=True, fast_bias_correction=False
)
总结¶
本教程介绍如何使用 torch.compile 和 OpenVINO 后端及 OpenVINO 量化器。有关 NNCF 和 NNCF 针对 PyTorch 模型的量化流程的更多细节，请参考 NNCF 量化指南 。有关更多信息，请查阅 通过 torch.compile 进行 OpenVINO 部署文档。
