PyTorch 2 导出训练后量化¶
作者: Jerry Zhang
本教程基于 torch._export.export，介绍了在图模式下进行训练后静态量化的步骤。与 FX 图模式量化 相比，这种流程预计将具有更高的模型覆盖率（在 14K 个模型上覆盖率为 88%）、更好的可编程性以及简化的用户体验。
可以使用 torch.export.export 导出的模型是使用此流程的前提，您可以在 Export DB 中找到支持的构造。
量化 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.6+, for more details on lower pytorch versions
# please check `Export the model with torch.export` section
m = torch.export.export(m, example_inputs).module()
# we get a model with aten ops
# Step 2. quantization
from torchao.quantization.pt2e.quantize_pt2e import (
prepare_pt2e,
convert_pt2e,
)
# install executorch: `pip install executorch`
from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
get_symmetric_quantization_config,
XNNPACKQuantizer,
)
# 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` 中，我们将 observer/ `fake_quant` observer 类暴露为一个对象供用户配置量化，这增加了用户可能需要关注的项。例如，不仅是 `dtype`，还有观察是如何发生的，这些都可以隐藏起来，使用户流程更简单。
这是新 API 的好处总结
可编程性（解决 1 和 2）：当用户的量化需求未被可用量化器覆盖时，用户可以构建自己的量化器，并如上所述将其与其他量化器组合。
简化的用户体验（解决 3）：提供一个后端和用户交互的单一实例。因此，您不再需要用户面对的量化配置映射用户意图，以及后端交互的单独量化配置来配置后端支持的内容。我们仍然会提供一个方法供用户查询量化器支持的内容。通过单一实例，组合不同的量化能力也变得比以前更自然。
例如，XNNPACK 不支持 `embedding_byte`，而我们在 ExecuTorch 中原生支持它。因此，如果我们有一个 `ExecuTorchQuantizer` 只量化 `embedding_byte`，那么它可以与 `XNNPACKQuantizer` 组合。（以前，这曾是将两个 `BackendConfig` 连接在一起，并且由于 `QConfigMapping` 中的选项不是后端特定的，用户还需要自己找出如何指定与组合后端量化能力匹配的配置。通过单一量化器实例，我们可以组合两个量化器并查询组合量化器的能力，这使得错误更少且更清晰，例如，`composed_quantizer.quantization_capabilities()`）。
关注点分离（解决 4）：在设计量化器 API 时，我们还将量化规范（以 `dtype`、最小/最大（位数）、对称等形式表示）与 observer 概念解耦。目前，observer 同时捕获了量化规范和如何观察（直方图 vs MinMax observer）。通过此更改，模型用户无需再与 observer 和 fake quant 对象进行交互。
定义辅助函数并准备数据集¶
我们将从进行必要的导入、定义一些辅助函数和准备数据开始。这些步骤与《PyTorch 中使用 Eager 模式进行静态量化》中的内容相同。
要使用整个 ImageNet 数据集运行本教程中的代码，请首先按照 ImageNet 数据 的说明下载 ImageNet。将下载的文件解压到 `data_path` 文件夹中。
下载 torchvision resnet18 模型 并将其重命名为 `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'torchao.quantization.pt2e'
)
# 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")
将模型设置为 eval 模式¶
对于训练后量化，我们需要将模型设置为 eval 模式。
model_to_quantize.eval()
使用 torch.export 导出模型¶
以下是如何使用 `torch.export` 导出模型
example_inputs = (torch.rand(2, 3, 224, 224),)
# for pytorch 2.6+
exported_model = torch.export.export(model_to_quantize, example_inputs).module()
# for pytorch 2.5 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.6+
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(model_to_quantize, example_inputs, dynamic_shapes=dynamic_shapes).module()
# for pytorch 2.5 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 executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
get_symmetric_quantization_config,
XNNPACKQuantizer,
)
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)
注意
请查看我们的 教程，其中介绍了如何编写新的 `Quantizer`。
准备模型以进行训练后量化¶
`prepare_pt2e` 将 `BatchNorm` 算子折叠到前面的 `Conv2d` 算子中，并在模型中适当的位置插入 observer。
prepared_model = prepare_pt2e(exported_model, quantizer)
print(prepared_model.graph)
校准¶
校准函数在模型中插入 observer 后运行。校准的目的是通过一些代表工作负载的样本示例（例如训练数据集的样本）进行运行，以便模型中的 observer 能够观察张量的统计信息，然后我们可以使用这些信息来计算量化参数。
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 表示（默认）
之前的文档中关于 表示，所有量化算子都表示为 `dequantize -> fp32_op -> qauntize`。
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
有关最新的参考表示，请参见 此处。
检查模型大小和准确率评估¶
现在我们可以将大小和模型准确率与基线模型进行比较。
# 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(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，它可以帮助在 eager 模式和 FX 图模式下进行调试。新版本的 Numeric Suite 与 PyTorch 2 导出模型配合使用仍在开发中。
降低和性能评估¶
此时生成的模型不是在设备上运行的最终模型，它是一个捕获用户预期量化计算的参考量化模型，表示为 ATen 算子和一些额外的量化/反量化算子。要获得在真实设备上运行的模型，我们需要降低模型。例如，对于在边缘设备上运行的模型，我们可以通过委派和 ExecuTorch 运行时算子进行降低。
结论¶
在本教程中，我们通过 `XNNPACKQuantizer` 详细介绍了 PyTorch 2 导出量化的整体量化流程，并获得了一个可以进一步降低到支持使用 XNNPACK 后端进行推理的后端的量化模型。要将其用于您自己的后端，请首先按照 教程 实现您后端的 `Quantizer`，然后使用该 `Quantizer` 量化模型。
