量化配方¶
Created On: Oct 26, 2020 | Last Updated: Mar 11, 2024 | Last Verified: Nov 05, 2024
本配方演示如何对PyTorch模型进行量化，以便以减小的尺寸和更快的推理速度运行，同时保持与原始模型大致相同的准确性。量化可应用于服务器和移动模型部署，但在移动端尤为重要甚至关键，因为非量化模型的大小可能超出iOS或Android应用程序允许的限制，导致部署或OTA更新耗时过长，并且推理速度过慢，影响用户体验。
简介¶
量化是一种技术，将模型参数中的32位浮点数转为8位整数。通过量化，模型大小和内存占用可减少到原来的1/4，并使推理速度提高约2-4倍，同时准确性几乎未受到影响。
总体来说，有三种方法或工作流程可以量化模型：训练后动态量化、训练后静态量化、量化感知训练。但是，如果您要使用的模型已经有量化版本，则可以直接使用，而无需经过以上三种工作流程。例如，`torchvision`库已经为MobileNet v2、ResNet 18、ResNet 50、Inception v3、GoogleNet等模型包含量化版本。因此，我们将最后一种方法视为一种简单的工作流程。
备注
量化支持的操作符集有限。请参阅`此处 <https://pytorch.org/blog/introduction-to-quantization-on-pytorch/#device-and-operator-support>`_了解更多信息。
工作流程¶
使用以下四种工作流程之一对模型进行量化。
1. 使用预训练的量化MobileNet v2模型¶
要获取MobileNet v2量化模型，只需执行：
import torchvision
model_quantized = torchvision.models.quantization.mobilenet_v2(pretrained=True, quantize=True)
比较未量化的MobileNet v2模型与其量化版本的大小差异：
model = torchvision.models.mobilenet_v2(pretrained=True)
import os
import torch
def print_model_size(mdl):
torch.save(mdl.state_dict(), "tmp.pt")
print("%.2f MB" %(os.path.getsize("tmp.pt")/1e6))
os.remove('tmp.pt')
print_model_size(model)
print_model_size(model_quantized)
输出将是：
14.27 MB
3.63 MB
2. 训练后动态量化¶
要应用动态量化，将模型中的所有权重从32位浮点数转换为8位整数，但只在计算激活值前将激活值转换为int8，只需调用`torch.quantization.quantize_dynamic`：
model_dynamic_quantized = torch.quantization.quantize_dynamic(
model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8
)
其中`qconfig_spec`指定`model`中要应用量化的子模块名称列表。
警告
动态量化的一个重要限制是，它目前仅支持`qconfig_spec`中的`nn.Linear`和`nn.LSTM`，这意味着要对其他模块（如`nn.Conv2d`）进行量化，需使用后面将讨论的静态量化或量化感知训练。
quantize_dynamic`API的完整文档在`此处。还可以参阅三个动态量化训练后示例：Bert示例，LSTM模型示例，以及另一个`LSTM模型示例 <https://pytorch.org/tutorials/recipes/recipes/dynamic_quantization.html#do-the-quantization>`_。
3. 训练后静态量化¶
此方法事先将权重和激活值都转换为8位整数，这样推理过程中激活值不需要动态转换。相比动态量化，训练后静态量化可以显著提升推理速度和减少模型大小，但可能对模型的原始准确性有更多影响。
要对模型进行静态量化，运行以下代码：
backend = "qnnpack"
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend
model_static_quantized = torch.quantization.prepare(model, inplace=False)
model_static_quantized = torch.quantization.convert(model_static_quantized, inplace=False)
运行`print_model_size(model_static_quantized)`后，可以看到静态量化后的模型大小为`3.98MB`。
完整的模型定义和静态量化示例在`此处 <https://pytorch.org/docs/stable/quantization.html#quantization-api-summary>`_。专门的静态量化教程在`此处 <https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html>`_。
备注
为了使模型能在通常具有ARM架构的移动设备上运行，需为`backend`使用`qnnpack`；在具有x86架构的计算机上运行时，需使用`x86``（旧的`fbgemm`仍可用，但推荐使用默认的`x86`）。
4. 量化感知训练¶
量化感知训练在模型训练过程中对所有权重和激活值插入伪量化，从而在推理时获得比训练后量化方法更高的精度。这通常用于CNN模型。
为启用模型的量化感知训练，在模型定义的`__init__`方法中定义一个`QuantStub`和一个`DeQuantStub`，以便在浮点数和量化类型之间转换张量：
self.quant = torch.quantization.QuantStub()
self.dequant = torch.quantization.DeQuantStub()
然后在模型定义的`forward`方法开始和结束时调用`x = self.quant(x)`和`x = self.dequant(x)`。
对量化感知训练，使用以下代码段：
model.qconfig = torch.quantization.get_default_qat_qconfig(backend)
model_qat = torch.quantization.prepare_qat(model, inplace=False)
# quantization aware training goes here
model_qat = torch.quantization.convert(model_qat.eval(), inplace=False)
关于量化感知训练的详细示例，可以参见`此处 <https://pytorch.org/docs/master/quantization.html#quantization-aware-training>`_和`此处 <https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#quantization-aware-training>`_。
预训练的量化模型也可以用于量化感知迁移学习，使用上文显示的相同`quant`和`dequant`调用。参见`此处 <https://pytorch.org/tutorials/intermediate/quantized_transfer_learning_tutorial.html#part-1-training-a-custom-classifier-based-on-a-quantized-feature-extractor>`_获取完整示例。
在使用上述步骤生成量化模型后，在模型能用于移动设备运行之前，需要进一步将其转换为`TorchScript`格式并优化以支持移动应用。详情参见`脚本和优化移动配方 <script_optimized.html>`_。
