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 部署文档。