适用于 OpenVINO torch.compile 后端的 PyTorch 2 导出量化

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

先决条件

• PyTorch 2 导出训练后量化

• 如何为 PyTorch 2 导出量化编写量化器

介绍

注意

这是一项实验性功能，量化 API 可能会发生变化。

本教程演示了如何在 PyTorch 2 导出量化流程中使用来自 神经网络压缩框架 (NNCF) 的 OpenVINOQuantizer，以生成针对 OpenVINO torch.compile 后端定制的量化模型，并解释如何将量化模型降低到 OpenVINO 表示。 OpenVINOQuantizer 通过专为 OpenVINO 设计的量化器放置，充分发挥了低精度 OpenVINO 内核的潜力。

PyTorch 2 导出量化流程使用 torch.export 将模型捕获到图中，并在 ATen 图之上执行量化转换。这种方法预计将具有显着更高的模型覆盖率、更高的灵活性和简化的用户体验。OpenVINO 后端将 TorchDynamo 生成的 FX 图编译为优化的 OpenVINO 模型。

量化流程主要包括四个步骤

• 步骤 1：基于 torch 导出机制，从即时模式（eager）模型中捕获 FX 图。

• 步骤 2：基于捕获的 FX 图，应用 PyTorch 2 导出量化流程与 OpenVINOQuantizer。

• 步骤 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 distribution 轻松安装。

pip install -U pip
pip install openvino, nncf

1. 捕获 FX 图

我们将首先执行必要的导入，从即时模式（eager）模块中捕获 FX 图。

import copy
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(), 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 的更多详细信息，请参阅 文档。

导入后端特定量化器后，我们将为后训练量化准备模型。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)

优化后的模型使用专为 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)
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 部署文档。