Pytorch量化(torch.quantization)

      在深度学习中，量化指的是使用更少的bit来存储原本以浮点数存储的tensor，以及使用更少的bit来完成原本以浮点数完成的计算。这么做的好处主要有如下几点：

• 更少的模型体积，接近4倍的减少；
• 可以更快的计算，由于更少的内存访问和更快的int8计算，可以快2~4倍。

      一个量化后的模型，其部分或者全部的tensor操作会使用int类型来计算，而不是使用量化之前的float类型。当然，量化还需要底层硬件支持，x86 CPU（支持AVX2）、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP这些主流硬件都对量化提

了支持。

      与典型的 FP32 型号相比，PyTorch 支持 INT8 量化，从而可将模型大小减少 4 倍，并将内存带宽要求减少 4 倍。 与 FP32 计算相比，对 INT8 计算的硬件支持通常快 2 到 4 倍。 量化主要是一种加速推理的技术，并且量化算子仅支持前向传递。PyTorch 支持多种方法来量化深度学习模型。 在大多数情况下，该模型在 FP32 中训练，然后将模型转换为 INT8

PyTorch对量化的支持目前有如下三种方式：

• Post Training Dynamic Quantization，模型训练完毕后的动态量化；
• Post Training Static Quantization，模型训练完毕后的静态量化；
• QAT（Quantization Aware Training），模型训练中开启量化。

示例: 模型量化后保存为jit形式

import torch
from torch import nn
from torch import quantization

class ConvBnReluModel(nn.Module):
    def __init__(self) -> None:
        super(ConvBnReluModel, self).__init__()
        self.conv = nn.Conv2d(3,5,3, bias=False)
        self.bn = nn.BatchNorm2d(5)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

m = ConvBnReluModel()
m.eval()
layers = [['conv','bn','relu']]
f = quantization.fuse_modules(m,layers, inplace=True)

types_to_quantize = {nn.Conv2d, nn.BatchNorm2d, nn.ReLU}
q = quantization.quantize_dynamic(f, types_to_quantize, dtype=torch.qint8)

s = torch.jit.script(q)
torch.jit.save(s, 'quantize_model.pth')