(原型) PyTorch 2 导出训练后量化¶
作者: Jerry Zhang
本教程介绍了基于 torch._export.export 在图模式下进行训练后静态量化的步骤。与 FX 图模式量化 相比,该流程预计将具有更高的模型覆盖率 (14K 模型的 88%)、更好的可编程性和更简单的 UX。
可通过 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
from torch._export import capture_pre_autograd_graph
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
# NOTE: this API will be updated to torch.export API in the future, but the captured
# result shoud mostly stay the same
m = capture_pre_autograd_graph(m, *example_inputs)
# 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 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
中,我们将观察者/fake_quant
观察者类作为对象公开,供用户配置量化,这增加了用户可能需要关注的事项。例如,不仅是dtype
,还有如何进行观察,这些可以从用户那里隐藏起来,以简化用户流程。
以下是新 API 的优势总结
可编程性(解决 1. 和 2.):当用户对量化的需求没有被现有的量化器覆盖时,用户可以构建自己的量化器,并将其与上述其他量化器组合。
简化的用户体验(解决 3.):提供一个实例,后端和用户都通过它进行交互。因此,您不再需要用户面对量化配置映射来映射用户的意图,以及一个单独的量化配置,后端与它交互以配置后端支持的内容。我们仍然会有一种方法供用户查询量化器中支持的内容。使用单个实例,组合不同的量化能力也比以前更自然。
例如,XNNPACK 不支持
embedding_byte
,而我们在 ExecuTorch 中对此有本地支持。因此,如果我们有ExecuTorchQuantizer
仅量化embedding_byte
,那么它可以与XNNPACKQuantizer
组合。 (以前,这曾经是将两个BackendConfig
连接起来,并且由于QConfigMapping
中的选项不是后端特定的,因此用户还需要弄清楚如何自行指定与组合后的后端的量化能力匹配的配置。使用单个量化器实例,我们可以组合两个量化器并查询组合后的量化器的功能,这使其不易出错且更清晰,例如,composed_quantizer.quantization_capabilities())
。关注点分离(解决 4.):在我们设计量化器 API 时,我们还将量化规范的规范从观察者概念中分离出来,量化规范以
dtype
、最小值/最大值(位数)、对称等形式表达。目前,观察者捕获了量化规范和如何观察(直方图与最小值最大值观察者)。通过这种改变,建模用户不再需要与观察者和伪量化对象交互。
定义辅助函数和准备数据集¶
我们将从进行必要的导入、定义一些辅助函数和准备数据开始。这些步骤与 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'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):
if isinstance(model, torch.jit.RecursiveScriptModule):
torch.jit.save(model, "temp.p")
else:
torch.jit.save(torch.jit.script(model), "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
导出模型
from torch._export import capture_pre_autograd_graph
example_inputs = (torch.rand(2, 3, 224, 224),)
exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs)
# or capture with dynamic dimensions
# from torch._export import dynamic_dim
# exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs, constraints=[dynamic_dim(example_inputs[0], 0)])
capture_pre_autograd_graph
是一个短期 API,将在官方 torch.export
API 就绪时更新为使用该 API。
导入后端特定量化器并配置如何量化模型¶
以下代码片段描述了如何量化模型
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
运算符中,并在模型中适当的位置插入观察者。
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 表示形式(默认)
之前关于 表示形式 的文档,所有量化运算符都表示为
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
参考量化模型表示形式(在 nightly 构建中可用)
我们将对选定运算符有特殊的表示形式,例如,量化线性。其他运算符表示为
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
scripted_float_model_file = "resnet18_scripted.pth"
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")
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
降低和性能评估¶
目前产生的模型不是在设备上运行的最终模型,它是一个参考量化模型,捕获了用户表达的意图的量化计算,以 ATen 运算符和一些额外的量化/反量化运算符的形式表达,为了获得在真实设备上运行的模型,我们需要降低模型。例如,对于在边缘设备上运行的模型,我们可以使用委托和 ExecuTorch 运行时运算符来降低模型。