（原型）PyTorch 2 导出后训练量化

Created On: Oct 02, 2023 | Last Updated: Oct 23, 2024 | Last Verified: Nov 05, 2024

作者：Jerry Zhang

本教程介绍了如何基于 torch._export.export 在图模式中执行后训练静态量化的步骤。相比 FX 图模式量化，此流程预计具有显著更高的模型覆盖率（88% 在 14K 模型中），更便于编程，以及简化的用户体验。

使用此流程的前提条件是支持 torch.export.export 导出功能，您可以在 导出数据库 中找到支持的结构。

具有量化器的量化流程 2 的高层架构如下所示：

float_model(Python)                          Example Input
    \                                              /
     \                                            /
—-------------------------------------------------------
|                        export                        |
—-------------------------------------------------------
                            |
                    FX Graph in ATen     Backend Specific Quantizer
                            |                       /
—--------------------------------------------------------
|                     prepare_pt2e                      |
—--------------------------------------------------------
                            |
                     Calibrate/Train
                            |
—--------------------------------------------------------
|                    convert_pt2e                       |
—--------------------------------------------------------
                            |
                    Quantized Model
                            |
—--------------------------------------------------------
|                       Lowering                        |
—--------------------------------------------------------
                            |
        Executorch, Inductor or <Other Backends>

PyTorch 2 导出量化 API 看起来如下所示：

import torch
class M(torch.nn.Module):
   def __init__(self):
      super().__init__()
      self.linear = torch.nn.Linear(5, 10)

   def forward(self, x):
      return self.linear(x)


example_inputs = (torch.randn(1, 5),)
m = M().eval()

# Step 1. program capture
# This is available for pytorch 2.5+, for more details on lower pytorch versions
# please check `Export the model with torch.export` section
m = torch.export.export_for_training(m, example_inputs).module()
# we get a model with aten ops


# Step 2. quantization
from torch.ao.quantization.quantize_pt2e import (
  prepare_pt2e,
  convert_pt2e,
)

from torch.ao.quantization.quantizer.xnnpack_quantizer import (
  XNNPACKQuantizer,
  get_symmetric_quantization_config,
)
# backend developer will write their own Quantizer and expose methods to allow
# users to express how they
# want the model to be quantized
quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
m = prepare_pt2e(m, quantizer)

# calibration omitted

m = convert_pt2e(m)
# we have a model with aten ops doing integer computations when possible

PyTorch 2 导出量化的动机

在 PyTorch 2 之前的版本中，我们有 FX 图模式量化，使用 QConfigMapping 和 BackendConfig 来进行定制。QConfigMapping 允许模型用户指定如何对模型进行量化，BackendConfig 允许后端开发者指定后端支持的量化方式。虽然该 API 在大多数情况下覆盖得相对较好，但它并不是完全可扩展的。当前 API 存在两个主要限制：

• 表达复杂操作模式的量化意图（如何观察/量化操作模式）的局限性：使用现有对象``QConfig`` 和 QConfigMapping 。

• 用户在表达他们希望如何对模型进行量化的意图时支持有限。例如，如果用户希望每隔一个线性层进行量化，或者量化行为与张量的实际形状有某种依赖关系（例如，当线性层具有3D输入时仅观察/量化输入和输出），后台开发者或建模用户需要更改核心量化API/流程。

一些改进可以使现有流程更好：

• 我们将“QConfigMapping”和“BackendConfig”作为独立的对象，“QConfigMapping”描述了用户希望如何对模型进行量化的意图，“BackendConfig”描述了后端支持哪种量化。“BackendConfig”是后端特定的，但“QConfigMapping”不是，用户可以提供与特定“BackendConfig”不兼容的“QConfigMapping”，这不是一个很好的用户体验。理想情况下，我们可以通过使配置(“QConfigMapping”)和量化能力(“BackendConfig”)都更具后端特性来更好地构造它们，从而减少关于不兼容的困惑。

• 在“QConfig”中，我们将观察者/“假量化”观察者类作为对象暴露给用户以配置量化，这增加了用户需要关心的事项。例如，不仅是“dtype”，还有观察应该如何进行，这些可以潜在地隐藏起来，使用户流程更简单。

以下是新API的优点摘要：

• **可编程性**（解决问题1和问题2）：当用户的量化需求未被现有量化器覆盖时，用户可以构建自己的量化器并像上面提到的那样与其他量化器组合。

• **简化用户体验**（解决问题3）：提供一个单一实例，与后端和用户交互。因此，用户不再需要面对量化配置映射以映射用户意图，以及与后端交互以配置后端支持的单独量化配置。我们仍然会有一种方法让用户查询量化器支持的内容。通过单一实例，组合不同的量化能力也比以前更自然。

  例如，XNNPACK不支持“embedding_byte”，而我们在ExecuTorch中本地支持它。因此，如果我们有一个只量化“embedding_byte”的“ExecuTorchQuantizer”，则可以与“XNNPACKQuantizer”组合。（之前，这通常是将两个“BackendConfig”连接在一起，由于“QConfigMapping”中的选项不是后端特定的，用户还需要自行弄清楚如何指定与组合后端的量化能力匹配的配置。通过单一量化器实例，我们可以组合两个量化器并查询组合量化器的能力，这使得它更少易出错且更清晰，例如，“composed_quantizer.quantization_capabilities()”。

• **关注点分离**（解决问题4）：在我们设计量化器API时，我们还将量化的规范与观察者概念解耦。当前，观察者既捕获了量化规范又捕获了如何观察（直方图vs最小最大观察者）。通过更改，建模用户不再与观察者和假量化对象交互。

定义辅助函数并准备数据集

我们将首先进行必要的导入，定义一些助手函数并准备数据。这些步骤与`PyTorch中的Eager模式静态量化 <https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html>`_中的步骤完全相同。

要使用整个ImageNet数据集运行本教程中的代码，请首先按照此处的说明下载ImageNet`ImageNet Data <http://www.image-net.org/download>`_。将下载的文件解压到``data_path``文件夹中。

下载`torchvision resnet18 模型 <https://download.pytorch.org/models/resnet18-f37072fd.pth>`_ 并将其重命名为``data/resnet18_pretrained_float.pth``。

import os
import sys
import time
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

import torchvision
from torchvision import datasets
from torchvision.models.resnet import resnet18
import torchvision.transforms as transforms

# Set up warnings
import warnings
warnings.filterwarnings(
    action='ignore',
    category=DeprecationWarning,
    module=r'.*'
)
warnings.filterwarnings(
    action='default',
    module=r'torch.ao.quantization'
)

# Specify random seed for repeatable results
_ = torch.manual_seed(191009)


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


def accuracy(output, target, topk=(1,)):
    """
    Computes the accuracy over the k top predictions for the specified
    values of k.
    """
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def evaluate(model, criterion, data_loader):
    model.eval()
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    cnt = 0
    with torch.no_grad():
        for image, target in data_loader:
            output = model(image)
            loss = criterion(output, target)
            cnt += 1
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            top1.update(acc1[0], image.size(0))
            top5.update(acc5[0], image.size(0))
    print('')

    return top1, top5

def load_model(model_file):
    model = resnet18(pretrained=False)
    state_dict = torch.load(model_file, weights_only=True)
    model.load_state_dict(state_dict)
    model.to("cpu")
    return model

def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print("Size (MB):", os.path.getsize("temp.p")/1e6)
    os.remove("temp.p")

def prepare_data_loaders(data_path):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    dataset = torchvision.datasets.ImageNet(
        data_path, split="train", transform=transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))
    dataset_test = torchvision.datasets.ImageNet(
        data_path, split="val", transform=transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ]))

    train_sampler = torch.utils.data.RandomSampler(dataset)
    test_sampler = torch.utils.data.SequentialSampler(dataset_test)

    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=train_batch_size,
        sampler=train_sampler)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=eval_batch_size,
        sampler=test_sampler)

    return data_loader, data_loader_test

data_path = '~/.data/imagenet'
saved_model_dir = 'data/'
float_model_file = 'resnet18_pretrained_float.pth'

train_batch_size = 30
eval_batch_size = 50

data_loader, data_loader_test = prepare_data_loaders(data_path)
example_inputs = (next(iter(data_loader))[0])
criterion = nn.CrossEntropyLoss()
float_model = load_model(saved_model_dir + float_model_file).to("cpu")
float_model.eval()

# create another instance of the model since
# we need to keep the original model around
model_to_quantize = load_model(saved_model_dir + float_model_file).to("cpu")

将模型设为评估模式

对于后训练量化，我们需要将模型设置为评估模式。

model_to_quantize.eval()

使用torch.export导出模型

以下是如何使用``torch.export``导出模型：

example_inputs = (torch.rand(2, 3, 224, 224),)
# for pytorch 2.5+
exported_model = torch.export.export_for_training(model_to_quantize, example_inputs).module()

# for pytorch 2.4 and before
# from torch._export import capture_pre_autograd_graph
# exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs)

# or capture with dynamic dimensions
# for pytorch 2.5+
dynamic_shapes = tuple(
  {0: torch.export.Dim("dim")} if i == 0 else None
  for i in range(len(example_inputs))
)
exported_model = torch.export.export_for_training(model_to_quantize, example_inputs, dynamic_shapes=dynamic_shapes).module()

# for pytorch 2.4 and before
# dynamic_shape API may vary as well
# from torch._export import dynamic_dim
# exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs, constraints=[dynamic_dim(example_inputs[0], 0)])

导入特定后端的量化器并配置如何对模型进行量化

以下代码片段描述了如何对模型进行量化：

from torch.ao.quantization.quantizer.xnnpack_quantizer import (
  XNNPACKQuantizer,
  get_symmetric_quantization_config,
)
quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config())

``Quantizer``是后端特定的，每个``Quantizer``将提供自己的方式来允许用户配置其模型。作为示例，以下是``XNNPackQuantizer``支持的不同配置API：

quantizer.set_global(qconfig_opt)  # qconfig_opt is an optional quantization config
    .set_object_type(torch.nn.Conv2d, qconfig_opt) # can be a module type
    .set_object_type(torch.nn.functional.linear, qconfig_opt) # or torch functional op
    .set_module_name("foo.bar", qconfig_opt)

备注

查看我们的`教程 <https://pytorch.org/tutorials/prototype/pt2e_quantizer.html>`_，描述如何编写新的``Quantizer``。

为后训练量化准备模型

``prepare_pt2e``将``BatchNorm``操作折叠到前面的``Conv2d``操作中，并在模型中的适当位置插入观察者。

prepared_model = prepare_pt2e(exported_model, quantizer)
print(prepared_model.graph)

校准

校准功能在观察者插入到模型中后运行。校准的目的是通过一些代表工作负载的示例（例如训练数据集的一个样本）运行，以便模型中的观察者能够观察张量的统计信息，我们可以稍后使用这些信息计算量化参数。

def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)
calibrate(prepared_model, data_loader_test)  # run calibration on sample data

将校准后的模型转换为量化模型

``convert_pt2e``接受一个校准后的模型并生成一个量化模型。

quantized_model = convert_pt2e(prepared_model)
print(quantized_model)

在这一步，我们目前有两种表示方式可供选择，但长期提供的确切表示可能根据PyTorch用户的反馈而有所改变。

• Q/DQ表示（默认）

  之前的文档`表示 <https://github.com/pytorch/rfcs/blob/master/RFC-0019-Extending-PyTorch-Quantization-to-Custom-Backends.md>`_中，所有量化操作都表示为``dequantize -> fp32_op -> quantize``。

def quantized_linear(x_int8, x_scale, x_zero_point, weight_int8, weight_scale, weight_zero_point, bias_fp32, output_scale, output_zero_point):
    x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor(
             x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8)
    weight_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor(
             weight_i8, weight_scale, weight_zero_point, weight_quant_min, weight_quant_max, torch.int8)
    weight_permuted = torch.ops.aten.permute_copy.default(weight_fp32, [1, 0]);
    out_fp32 = torch.ops.aten.addmm.default(bias_fp32, x_fp32, weight_permuted)
    out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor(
    out_fp32, out_scale, out_zero_point, out_quant_min, out_quant_max, torch.int8)
    return out_i8

• 参考量化模型表示

  我们将对选定的操作（例如量化线性）提供特殊表示。其他操作表示为``dq -> float32_op -> q``且``q/dq``分解为更原始的操作。您可以通过使用``convert_pt2e(…, use_reference_representation=True)``获得这种表示。

# Reference Quantized Pattern for quantized linear
def quantized_linear(x_int8, x_scale, x_zero_point, weight_int8, weight_scale, weight_zero_point, bias_fp32, output_scale, output_zero_point):
    x_int16 = x_int8.to(torch.int16)
    weight_int16 = weight_int8.to(torch.int16)
    acc_int32 = torch.ops.out_dtype(torch.mm, torch.int32, (x_int16 - x_zero_point), (weight_int16 - weight_zero_point))
    bias_scale = x_scale * weight_scale
    bias_int32 = out_dtype(torch.ops.aten.div.Tensor, torch.int32, bias_fp32, bias_scale)
    acc_int32 = acc_int32 + bias_int32
    acc_int32 = torch.ops.out_dtype(torch.ops.aten.mul.Scalar, torch.int32, acc_int32, x_scale * weight_scale / output_scale) + output_zero_point
    out_int8 = torch.ops.aten.clamp(acc_int32, qmin, qmax).to(torch.int8)
    return out_int8

查看`这里 <https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py>`_以获取最新的参考表示。

检查模型大小和准确性评估

现在我们可以将模型的大小和准确性与基准模型进行比较。

# Baseline model size and accuracy
print("Size of baseline model")
print_size_of_model(float_model)

top1, top5 = evaluate(float_model, criterion, data_loader_test)
print("Baseline Float Model Evaluation accuracy: %2.2f, %2.2f"%(top1.avg, top5.avg))

# Quantized model size and accuracy
print("Size of model after quantization")
# export again to remove unused weights
quantized_model = torch.export.export_for_training(quantized_model, example_inputs).module()
print_size_of_model(quantized_model)

top1, top5 = evaluate(quantized_model, criterion, data_loader_test)
print("[before serilaization] Evaluation accuracy on test dataset: %2.2f, %2.2f"%(top1.avg, top5.avg))

备注

由于模型尚未降低到目标设备，我们目前无法进行性能评估，它仅是ATen操作中量化计算的表示。

备注

权重现在仍是fp32，我们以后可能会对量化操作进行常量传播以获取整数权重。

如果您希望获得更好的准确性或性能，请尝试以不同方式配置``quantizer``，每个``quantizer``都会有自己的配置方式，因此请查阅您正在使用的量化器的文档以了解更多关于如何控制模型量化的信息。

保存和加载量化模型

我们将展示如何保存和加载量化模型。

# 0. Store reference output, for example, inputs, and check evaluation accuracy:
example_inputs = (next(iter(data_loader))[0],)
ref = quantized_model(*example_inputs)
top1, top5 = evaluate(quantized_model, criterion, data_loader_test)
print("[before serialization] Evaluation accuracy on test dataset: %2.2f, %2.2f"%(top1.avg, top5.avg))

# 1. Export the model and Save ExportedProgram
pt2e_quantized_model_file_path = saved_model_dir + "resnet18_pt2e_quantized.pth"
# capture the model to get an ExportedProgram
quantized_ep = torch.export.export(quantized_model, example_inputs)
# use torch.export.save to save an ExportedProgram
torch.export.save(quantized_ep, pt2e_quantized_model_file_path)


# 2. Load the saved ExportedProgram
loaded_quantized_ep = torch.export.load(pt2e_quantized_model_file_path)
loaded_quantized_model = loaded_quantized_ep.module()

# 3. Check results for example inputs and check evaluation accuracy again:
res = loaded_quantized_model(*example_inputs)
print("diff:", ref - res)

top1, top5 = evaluate(loaded_quantized_model, criterion, data_loader_test)
print("[after serialization/deserialization] Evaluation accuracy on test dataset: %2.2f, %2.2f"%(top1.avg, top5.avg))

输出：

[before serialization] Evaluation accuracy on test dataset: 79.82, 94.55
diff: tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]])

[after serialization/deserialization] Evaluation accuracy on test dataset: 79.82, 94.55

调试量化模型

您可以使用`Numeric Suite <https://pytorch.org/docs/stable/quantization-accuracy-debugging.html#numerical-debugging-tooling-prototype>`_，它可以帮助以急切模式和FX图模式进行调试。与PyTorch 2 Export模型一起工作的新版本Numeric Suite仍在开发中。

降低和性能评估

目前生成的模型不是在设备上运行的最终模型，它是一个参考量化模型，捕获用户意图的量化计算，表示为ATen操作和一些额外的量化/去量化操作。为了获得在真实设备上运行的模型，我们需要降低模型。例如，对于在边缘设备上运行的模型，我们可以使用委托和ExecuTorch运行时操作进行降低。

总结

在本教程中，我们使用``XNNPACKQuantizer``完成了PyTorch 2 Export量化的整体量化流程并获得了一个可以进一步降低到支持XNNPACK后端推理的后端的量化模型。要将其用于您自己的后端，请首先查看`教程 <https://pytorch.org/tutorials/prototype/pt2e_quantizer.html>`__并为您的后端实现一个``Quantizer``，然后使用该``Quantizer``进行模型量化。