unstable/openvino_quantizer

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

面向 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 表示。由于放置了专为 OpenVINO 设计的量化器，OpenVINOQuantizer 释放了低精度 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 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 的更多详细信息，请参阅文档。

在我们导入特定于后端的量化器后，我们将为训练后量化准备模型。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 图可以利用 OpenVINO 优化，通过 torch.compile(…, backend=”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 设计的低级内核。与即时模式（eager）模型相比，这应能显著加快推理时间。

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