(prototype) PyTorch 2 Export 训练后量化¶
创建于: 2023 年 10 月 02 日 | 最后更新于: 2024 年 10 月 23 日 | 最后验证于: 2024 年 11 月 05 日
作者: Jerry Zhang
本教程介绍基于 torch._export.export 在图模式下进行训练后静态量化的步骤。与 FX 图模式量化 相比，此流程预计将具有显著更高的模型覆盖率（在 1.4 万个模型上达到 88%）、更好的可编程性以及简化的用户体验 (UX)。
torch.export.export 可导出是使用此流程的先决条件，您可以在 Export DB 中找到支持的构造。
使用 quantizer 的量化 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 export 量化 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 Export 量化的动机¶
在 PyTorch 2 之前的版本中，我们有 FX 图模式量化，它使用 QConfigMapping 和 BackendConfig 进行自定义。QConfigMapping 允许建模用户指定他们希望如何量化模型，BackendConfig 允许后端开发者指定他们的后端支持的量化方式。虽然该 API 相对较好地涵盖了大多数用例，但它不是完全可扩展的。当前 API 有两个主要限制：
在使用现有对象（
QConfig和QConfigMapping）表达复杂操作模式的量化意图（如何观察/量化操作模式）方面的限制。对用户如何表达他们希望如何量化模型的意图的支持有限。例如，如果用户想量化模型中隔一个线性层进行量化，或者量化行为依赖于张量的实际形状（例如，只有当线性层输入为 3D 时才观察/量化输入和输出），后端开发者或建模用户需要修改核心量化 API/流程。
一些改进可以使现有流程更好
我们将
QConfigMapping和BackendConfig作为独立的对象使用，QConfigMapping描述用户希望如何量化模型的意图，BackendConfig描述后端支持的量化类型。BackendConfig是后端特定的，但QConfigMapping不是，用户可以提供与特定BackendConfig不兼容的QConfigMapping，这并不是很好的用户体验 (UX)。理想情况下，我们可以通过使配置 (QConfigMapping) 和量化能力 (BackendConfig) 都成为后端特定的来更好地构建它，这样关于不兼容性的混淆就会减少。在
QConfig中，我们将 observer/fake_quantobserver 类作为对象暴露给用户进行量化配置，这增加了用户可能需要关注的事项。例如，不仅是dtype，还有观察应该如何发生，这些可以潜在地对用户隐藏，从而简化用户流程。
以下是新 API 的优势总结
可编程性（解决第 1 和 2 点）：当用户现有的量化需求不能由可用的量化器覆盖时，用户可以构建自己的量化器并将其与其他量化器组合，如上所述。
简化的用户体验 (UX)（解决第 3 点）：提供一个单一实例，后端和用户都可以与之交互。因此，您不再需要用户端的量化配置映射来映射用户意图，也不再需要独立的量化配置供后端交互以配置后端支持的内容。我们仍然会提供一个方法供用户查询量化器中支持的内容。通过单一实例，组合不同的量化能力也比以前更加自然，例如，
composed_quantizer.quantization_capabilities()，这使得错误更少且更简洁。例如，XNNPACK 不支持
embedding_byte，而 ExecuTorch 中 natively 支持。因此，如果有一个ExecuTorchQuantizer只量化embedding_byte，那么它可以与XNNPACKQuantizer组合。(以前，这需要将两个BackendConfig连接起来，并且由于QConfigMapping中的选项不是后端特定的，用户还需要自己找出如何指定与组合后端量化能力匹配的配置。通过单一的 quantizer 实例，我们可以组合两个 quantizer 并查询组合后的 quantizer 的能力，这使得错误更少且更简洁，例如，composed_quantizer.quantization_capabilities())。关注点分离（解决第 4 点）：在设计量化器 API 时，我们还将量化规范（以
dtype、min/max（位数）、对称性等表示）与 observer 概念解耦。目前，observer 同时捕获量化规范和如何观察（Histogram vs MinMax observer）。通过这一改变，建模用户无需与 observer 和 fake quant 对象交互。
定义辅助函数并准备数据集¶
我们首先进行必要的导入，定义一些辅助函数并准备数据。这些步骤与 PyTorch Eager Mode 静态量化完全相同。
要使用整个 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'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")
使用 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)])
导入特定于后端的 Quantizer 并配置如何量化模型¶
以下代码片段描述了如何量化模型
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)
注意
请查阅我们的 教程，它描述了如何编写新的 Quantizer。
准备模型进行训练后量化¶
prepare_pt2e 将 BatchNorm 操作折叠到前面的 Conv2d 操作中，并在模型中适当的位置插入 observer。
prepared_model = prepare_pt2e(exported_model, quantizer)
print(prepared_model.graph)
校准（Calibration）¶
校准函数在 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_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 来帮助在 Eager Mode 和 FX 图模式下进行调试。与 PyTorch 2 Export 模型配合使用的新版 Numeric Suite 仍在开发中。
下沉与性能评估¶
此时生成的模型不是在设备上运行的最终模型，它是一个参考量化模型，用于捕获用户意图的量化计算，以 ATen 操作符和一些额外的量化/反量化操作符的形式表示，要获得在真实设备上运行的模型，我们需要对模型进行下沉。例如，对于在边缘设备上运行的模型，我们可以使用委托（delegation）和 ExecuTorch 运行时操作符进行下沉。
