量化菜谱
导航
量化菜谱#
参考：量化菜谱
量化是一种将模型参数中的 32 位浮点数转换为 8 位整数的技术。通过量化，模型大小和内存占用可以减少到原来大小的 1/4，推理速度可以提高 2-4 倍，而精度保持不变。
量化模型的方法或工作流程大致有三种：后训练动态量化（post training dynamic quantization）、后训练静态量化（post training dynamic quantization）和量化感知训练（quantization aware training）。但是，如果您想要使用的模型已经有量化版本，您可以直接使用它，而不需要经过上面三个工作流中的任何一个。例如，torchvision 库已经包括了模型 MobileNet v2、ResNet 18、ResNet 50、Inception v3、GoogleNet 等的量化版本。因此，我们将最后一种方法作为另一个工作流，尽管很简单。
备注
量化支持可用于有限的一组算子。
PyTorch 支持四种量化工作流。
使用预先训练的量化网络#
要得到网络的量化模型，只需做：
from torchvision.models.quantization import resnet18 as qresnet18
model_quantized = qresnet18(pretrained=True,
quantize=True)
为了比较非量化的模型与其量化版本的大小差异：
from torchvision.models import resnet18
model = resnet18(pretrained=True)
import os
import torch
def print_model_size(mdl, name=''):
torch.save(mdl.state_dict(), "tmp.pt")
size = os.path.getsize("tmp.pt")/1e6
print(f'{name} {size: .2f} MB')
os.remove('tmp.pt')
print_model_size(model, '原始模型')
print_model_size(model_quantized, '量化模型')
原始模型 46.83 MB
量化模型 11.84 MB
后训练动态量化#
要应用动态量化，它将模型中的所有权重从 32 位的浮点数转换为 8 位的整数，但在对这些激活执行计算之前不会将激活转换为 int8，只需调用 quantize_dynamic()：
from torch.quantization.quantize import quantize_dynamic
model_dynamic_quantized = quantize_dynamic(
model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8
)
其中 qconfig_spec 指定要应用量化的 model 中的子模块名称列表。
警告
动态量化的一个重要限制是，虽然它是最简单的工作流程，如果你没有预先训练的量化模型准备使用，它目前在 qconfig_spec 中只支持 nn.Linear 和 nn.LSTM，这意味着你将不得不使用静态量化或量化感知训练，稍后讨论，以量化其他模块，如 nn.Conv2d。
使用训练后动态量化的另外三个例子是 Bert 例子、LSTM 模型例子 和另一个 LSTM demo 例子。
后训练静态量化#
该方法将权值和激活提前转换为 8 位整数，这样就不会像动态量化那样在推理过程中对激活进行实时转换，从而显著提高了性能。
要在模型上应用静态量化，运行以下代码：
from torch.quantization.quantize import prepare, convert
from torch.quantization.qconfig import get_default_qconfig
from torchvision.models.quantization import resnet18 as qresnet18
backend = "fbgemm" # 若为 x86，否则为 'qnnpack'
model.qconfig = get_default_qconfig(backend)
torch.backends.quantized.engine = backend
model_static_quantized = prepare(model, inplace=False)
model_static_quantized = convert(model_static_quantized, inplace=False)
/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/torch/ao/quantization/observer.py:177: UserWarning: Please use quant_min and quant_max to specify the range for observers. reduce_range will be deprecated in a future release of PyTorch.
warnings.warn(
/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/torch/ao/quantization/observer.py:1124: UserWarning: must run observer before calling calculate_qparams. Returning default scale and zero point
warnings.warn(
显示静态量化模型大小：
print_model_size(model_static_quantized)
11.93 MB
这里有一个完整的模型定义和静态量化的例子。这里有一个专门的静态量化教程。
备注
为了使模型运行在通常有 ARM 架构的移动设备上，你需要使用 'qnnpack' 作为后台；要在 x86 架构的计算机上运行模型，请使用 'fbgemm'。
量化感知训练（QAT）#
量化感知训练（Quantization aware training）在模型训练过程中对所有的权值和激活量都插入伪量化，比训练后的量化方法具有更高的推理精度。它通常用于 CNN 的模型中。
要启用量化感知训练的模型，请在模型定义的 __init__ 方法中定义 QuantStub 和 DeQuantStub，以将张量从浮点类型转换为量化类型，反之亦然。
from torch import nn
from torch.quantization.stubs import QuantStub, DeQuantStub
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
# 其他运算
x = self.dequant(x)
return x
然后在模型定义的前向方法的开头和结尾分别调用 x = self.quant(x) 和 x = self.dequant(x)。
要进行量化感知训练，请使用下面的代码片段:
from torch.quantization.quantize import prepare_qat, convert
from torch.quantization.qconfig import get_default_qat_qconfig
model = MyModel()
model.qconfig = get_default_qat_qconfig(backend)
model_qat = prepare_qat(model, inplace=False)
# QAT
model_qat = convert(model_qat.eval(), inplace=False)
/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/torch/ao/quantization/utils.py:210: UserWarning: must run observer before calling calculate_qparams. Returning default values.
warnings.warn(
在使用上面的步骤之一生成量化模型之后，在模型可以在移动设备上运行之前，它需要进一步转换为 TorchScript 格式，然后针对移动端应用程序进行优化。请参阅脚本和优化移动端食谱的详细信息。