如何使用PyTorch的量化功能？

冷不防 2022-10-22 11:48 455阅读 0赞

> 点击上方“**机器学习与生成对抗网络**”，关注星标
> 
> 获取有趣、好玩的前沿干货！

来源 | PaperWeekly  作者 | Gemfield

**01. 背景**

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

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

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

PyTorch 1.1 的时候开始添加 torch.qint8 dtype、torch.quantize\_linear 转换函数来开始对量化提供有限的实验性支持。PyTorch 1.3 开始正式支持量化，在可量化的 Tensor 之外，PyTorch 开始支持 CNN 中最常见的 operator 的量化操作，包括：

1. Tensor 上的函数: view, clone, resize, slice, add, multiply, cat, mean, max, sort, topk；

2. 常见的模块（在 torch.nn.quantized 中）：Conv2d, Linear, Avgpool2d, AdaptiveAvgpool2d, MaxPool2d, AdaptiveMaxPool2d, Interpolate, Upsample；

3. 为了量化后还维持更高准确率的合并操作（在torch.nn.intrinsic中）：ConvReLU2d, ConvBnReLU2d, ConvBn2d，LinearReLU，add\_relu。

在 PyTorch 1.4 的时候，PyTorch 添加了 nn.quantized.Conv3d，与此同时，torchvision 0.5 开始提供量化版本的 ResNet、ResNext、MobileNetV2、GoogleNet、InceptionV3 和 ShuffleNetV2。

到 PyTorch 1.5 的时候，QNNPACK 添加了对 dynamic quantization 的支持，也就为量化版的 LSTM 在手机平台上使用提供了支撑——也就是添加了对 PyTorch mobile 的 dynamic quantization 的支持；增加了量化版本的 sigmoid、leaky relu、batch\_norm、BatchNorm2d、 Avgpool3d、quantized\_hardtanh、quantized ELU activation、quantized Upsample3d、quantized batch\_norm3d、 batch\_norm3d + relu operators的fused、quantized hardsigmoid。

在 PyTorch 1.6 的时候，添加了 quantized Conv1d、quantized hardswish、quantized layernorm、quantized groupnorm、quantized instancenorm、quantized reflection\_pad1d、quantized adaptive avgpool、quantized channel shuffle op、Quantized Threshold；添加 ConvBn3d, ConvBnReLU3d, BNReLU2d, BNReLU3d；per-channel 的量化得到增强；添加对 LSTMCell、RNNCell、GRUCell 的 Dynamic quantization 支持；在 nn.DataParallel 和  nn.DistributedDataParallel 中可以使用 Quantization aware training；支持 CUDA 上的 quantized tensor。

到目前的最新版本的 PyTorch 1.7，又添加了 Embedding 和 EmbeddingBag quantization、aten::repeat、aten::apend、tensor 的 stack、tensor 的 fill\_、per channel affine quantized tensor 的 clone、1D batch normalization、N-Dimensional constant padding、CELU operator、FP16 quantization 的支持。

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

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

在开始这三部分之前，先介绍下最基础的 Tensor 的量化。

##     ##

**02. Tensor的量化**

PyTorch 为了实现量化，首先就得需要具备能够表示量化数据的 Tensor，这就是从 PyTorch 1.1 之后引入的 Quantized Tensor。Quantized Tensor 可以存储  int8/uint8/int32 类型的数据，并携带有 scale、zero\_point 这些参数。把一个标准的 float Tensor 转换为量化 Tensor 的步骤如下：

>>> x = torch.rand(2,3, dtype=torch.float32) 
    >>> x
    tensor([[0.6839, 0.4741, 0.7451],
            [0.9301, 0.1742, 0.6835]])
    
    >>> xq = torch.quantize_per_tensor(x, scale = 0.5, zero_point = 8, dtype=torch.quint8)
    tensor([[0.5000, 0.5000, 0.5000],
            [1.0000, 0.0000, 0.5000]], size=(2, 3), dtype=torch.quint8,
           quantization_scheme=torch.per_tensor_affine, scale=0.5, zero_point=8)
    
    >>> xq.int_repr()
    tensor([[ 9,  9,  9],
            [10,  8,  9]], dtype=torch.uint8)

quantize\_per\_tensor 函数就是使用给定的 scale 和 zp 来把一个 float tensor 转化为quantized tensor，后文你还会遇到这个函数。通过上面这几个数的变化，你可以感受到，量化 tensor，也就是 xq，和 fp32 tensor 的关系大概就是:

xq = round(x / scale + zero_point)

scale 这个缩放因子和 zero\_point 是两个参数，建立起了 fp32 tensor 到量化 tensor 的映射关系。scale 体现了映射中的比例关系，而 zero\_point 则是零基准，也就是 fp32 中的零在量化 tensor 中的值。因为当 x 为零的时候，上述 xq 就变成了：

xq = round(zero_point) = zero_point

现在 xq 已经是一个量化 tensor 了，我们可以把 xq 在反量化回来，如下所示：

# xq is a quantized tensor with data represented as quint8
    >>> xdq = xq.dequantize()
    >>> xdq
    tensor([[0.5000, 0.5000, 0.5000],
            [1.0000, 0.0000, 0.5000]])

dequantize 函数就是 quantize\_per\_tensor 的反义词，把一个量化 tensor 转换为 float tensor。也就是：

xdq = (xq - zero_point) * scale

xdq 和 x 的值已经出现了偏差的事实告诉了我们两个道理：

*  量化会有精度损失；
 *  我们这里随便选取的 scale 和 zp 太烂，选择合适的 scale 和 zp 可以有效降低精度损失。不信你把 scale 和 zp 分别换成 scale = 0.0036, zero\_point = 0试试。

而在 PyTorch 中，选择合适的 scale 和 zp 的工作就由各种 observer 来完成。

Tensor 的量化支持两种模式：per tensor 和 per channel。Per tensor 是说一个 tensor 里的所有 value 按照同一种方式去 scale 和 offset；per channel 是对于  tensor 的某一个维度（通常是 channel 的维度）上的值按照一种方式去 scale 和 offset，也就是一个 tensor 里有多种不同的 scale 和 offset 的方式（组成一个vector），如此以来，在量化的时候相比 per tensor 的方式会引入更少的错误。PyTorch 目前支持 conv2d()、conv3d()、linear() 的 per channel 量化。

**03. Post Training Dynamic Quantization**

这种量化方式经常缩略前面的两个单词从而称之为 Dynamic Quantization，中文为动态量化。这是什么意思呢？你看到全称中的两个关键字了吗：Post、Dynamic：

*  Post：也就是训练完成后再量化模型的权重参数；
 *  Dynamic：也就是网络在前向推理的时候动态的量化 float32 类型的输入。

Dynamic Quantization 使用下面的 API 来完成模型的量化：

torch.quantization.quantize_dynamic(model, qconfig_spec=None, dtype=torch.qint8, mapping=None, inplace=False)

quantize\_dynamic 这个 API 把一个 float model 转换为 dynamic quantized model，也就是只有权重被量化的 model，dtype 参数可以取值 float16 或者  qint8。当对整个模型进行转换时，默认只对以下的 op 进行转换：

*  Linear
 *  LSTM
 *  LSTMCell
 *  RNNCell
 *  GRUCell

为啥呢？因为 dynamic quantization只是把权重参数进行量化，而这些 layer 一般参数数量很大，在整个模型中参数量占比极高，因此边际效益高。对其它 layer进行 dynamic quantization 几乎没有实际的意义。

再来说说这个 API 的第二个参数：qconfig\_spec：

*  qconfig\_spec 指定了一组 qconfig，具体就是哪个 op 对应哪个 qconfig；
 *  每个 qconfig 是 QConfig 类的实例，封装了两个 observer；
 *  这两个 observer 分别是 activation 的 observer 和 weight 的 observer；
 *  但是动态量化使用的是 QConfig 子类 QConfigDynamic 的实例，该实例实际上只封装了 weight 的 observer；
 *  activate 就是 post process，就是 op forward 之后的后处理，但在动态量化中不包含；
 *  observer 用来根据四元组（min\_val，max\_val，qmin, qmax）来计算 2 个量化的参数：scale 和 zero\_point；
 *  qmin、qmax 是算法提前确定好的，min\_val 和 max\_val 是从输入数据中观察到的，所以起名叫 observer。

当 qconfig\_spec 为 None 的时候就是默认行为，如果想要改变默认行为，则可以：

*  qconfig\_spec 赋值为一个 set，比如：\{nn.LSTM, nn.Linear\}，意思是指定当前模型中的哪些 layer 要被 dynamic quantization；
 *  qconfig\_spec 赋值为一个 dict，key 为 submodule 的 name 或 type，value 为 QConfigDynamic 实例（其包含了特定的 Observer，比如 MinMaxObserver、MovingAverageMinMaxObserver、PerChannelMinMaxObserver、MovingAveragePerChannelMinMaxObserver、HistogramObserver）。

事实上，当 qconfig\_spec 为 None 的时候，quantize\_dynamic API 就会使用如下的默认值：

qconfig_spec = {
                    nn.Linear : default_dynamic_qconfig,
                    nn.LSTM : default_dynamic_qconfig,
                    nn.GRU : default_dynamic_qconfig,
                    nn.LSTMCell : default_dynamic_qconfig,
                    nn.RNNCell : default_dynamic_qconfig,
                    nn.GRUCell : default_dynamic_qconfig,
                }

这就是 Gemfield 刚才提到的动态量化只量化 Linear 和 RNN 变种的真相。而 default\_dynamic\_qconfig 是 QConfigDynamic 的一个实例，使用如下的参数进行构造：

default_dynamic_qconfig = QConfigDynamic(activation=default_dynamic_quant_observer, weight=default_weight_observer)
    default_dynamic_quant_observer = PlaceholderObserver.with_args(dtype=torch.float, compute_dtype=torch.quint8)
    default_weight_observer = MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)

其中，用于 activation 的 PlaceholderObserver 就是个占位符，啥也不做；而用于 weight 的 MinMaxObserver 就是记录输入 tensor 中的最大值和最小值，用来计算 scale 和 zp。

对于一个默认行为下的 quantize\_dynamic 调用，你的模型会经历什么变化呢？Gemfield 使用一个小网络来演示下：

class CivilNet(nn.Module):
        def __init__(self):
            super(CivilNet, self).__init__()
            gemfieldin = 1
            gemfieldout = 1
            self.conv = nn.Conv2d(gemfieldin, gemfieldout, kernel_size=1, stride=1, padding=0, groups=1, bias=False)
            self.fc = nn.Linear(3, 2,bias=False)
            self.relu = nn.ReLU(inplace=False)
    
        def forward(self, x):
            x = self.conv(x)
            x = self.fc(x)
            x = self.relu(x)
            return x

原始网络和动态量化后的网络如下所示：

#原始网络
    CivilNet(
      (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (fc): Linear(in_features=3, out_features=2, bias=False)
      (relu): ReLU()
    )
    
    #quantize_dynamic 后
    CivilNet(
      (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (fc): DynamicQuantizedLinear(in_features=3, out_features=2, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
      (relu): ReLU()
    )

可以看到，除了 Linear，其它 op 都没有变动。而 Linear 被转换成了 DynamicQuantizedLinear，DynamicQuantizedLinear 就是 torch.nn.quantized.dynamic.modules.linear.Linear 类。

没错，quantize\_dynamic API 的本质就是检索模型中 op 的 type，如果某个 op 的 type 属于字典 DEFAULT\_DYNAMIC\_QUANT\_MODULE\_MAPPINGS 的 key，那么，这个 op 将被替换为 key 对应的 value：

# Default map for swapping dynamic modules
    DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS = {
        nn.GRUCell: nnqd.GRUCell,
        nn.Linear: nnqd.Linear,
        nn.LSTM: nnqd.LSTM,
        nn.LSTMCell: nnqd.LSTMCell,
        nn.RNNCell: nnqd.RNNCell,
    }

这里，nnqd.Linear 就是 DynamicQuantizedLinear 就是 torch.nn.quantized.dynamic.modules.linear.Linear。但是，type从key 换为 value，那这个新的 type 如何实例化呢？更重要的是，实例化新的 type 一定是要用之前的权重参数的呀。没错，以 Linear 为例，该逻辑定义在 nnqd.Linear 的 from\_float() 方法中，通过如下方式实例化：

new_mod = mapping[type(mod)].from_float(mod)

from\_float 做的事情主要就是：

*  使用 MinMaxObserver 计算模型中 op **权重**参数中 tensor 的最大值最小值（这个例子中只有 Linear op），缩小量化时原始值的取值范围，提高量化的精度；
 *  通过上述步骤中得到四元组中的 min\_val 和 max\_val，再结合算法确定的 qmin, qmax 计算出 scale 和 zp，参考前文“Tensor的量化”小节，计算得到量化后的weight，这个量化过程有torch.quantize\_per\_tensor 和 torch.quantize\_per\_channel两种，默认是前者（因为qchema默认是torch.per\_tensor\_affine）；
 *  实例化 nnqd.Linear，然后使用 qlinear.set\_weight\_bias 将量化后的 weight 和原始的 bias 设置到新的 layer 上。其中最后一步还涉及到 weight 和 bias 的打包，在源代码中是这样的：

#ifdef USE_FBGEMM
        if (ctx.qEngine() == at::QEngine::FBGEMM) {
          return PackedLinearWeight::prepack(std::move(weight), std::move(bias));
        }
    #endif
    
    #ifdef USE_PYTORCH_QNNPACK
        if (ctx.qEngine() == at::QEngine::QNNPACK) {
          return PackedLinearWeightsQnnp::prepack(std::move(weight), std::move(bias));
        }
    #endif
        TORCH_CHECK(false,"Didn't find engine for operation quantized::linear_prepack ",toString(ctx.qEngine()));

也就是说依赖 FBGEMM、QNNPACK 这些 backend。量化完后的模型在推理的时候有什么不一样的呢？在原始网络中，从输入到最终输出是这么计算的：

#input
    torch.Tensor([[[[-1,-2,-3],[1,2,3]]]])
    
    #经过卷积后（权重为torch.Tensor([[[[-0.7867]]]])）
    torch.Tensor([[[[ 0.7867,  1.5734,  2.3601],[-0.7867, -1.5734, -2.3601]]]])
    
    #经过fc后（权重为torch.Tensor([[ 0.4097, -0.2896, -0.4931], [-0.3738, -0.5541,  0.3243]]) )
    torch.Tensor([[[[-1.2972, -0.4004], [1.2972,  0.4004]]]])
    
    #经过relu后
    torch.Tensor([[[[0.0000, 0.0000],[1.2972, 0.4004]]]])

而在动态量化模型中，上述过程就变成了：

#input
    torch.Tensor([[[[-1,-2,-3],[1,2,3]]]])
    
    #经过卷积后（权重为torch.Tensor([[[[-0.7867]]]])）
    torch.Tensor([[[[ 0.7867,  1.5734,  2.3601],[-0.7867, -1.5734, -2.3601]]]])
    
    #经过fc后（权重为torch.Tensor([[ 0.4085, -0.2912, -0.4911],[-0.3737, -0.5563,  0.3259]], dtype=torch.qint8,scale=0.0043458822183310986,zero_point=0) )
    torch.Tensor([[[[-1.3038, -0.3847], [1.2856,  0.3969]]]])
    
    #经过relu后
    torch.Tensor([[[[0.0000, 0.0000], [1.2856, 0.3969]]]])

所以关键点就是这里的 Linear op 了，因为其它 op 和量化之前是一模一样的。你可以看到 Linear 权重的 scale 为 0.0043458822183310986，zero\_point 为0。scale 和 zero\_point 怎么来的呢？由其使用的 observer 计算得到的，具体来说就是默认的 MinMaxObserver，它是怎么工作的呢？还记得前面说过的 observer 负责根据四元组来计算 scale 和 zp 吧：

**在各种 observer 中，计算权重的 scale 和 zp 离不开这四个变量：min\_val，max\_val，qmin, qmax，分别代表 op 权重数据 /input tensor 数据分布的最小值和最大值，以及量化后的取值范围的最小、最大值。**

qmin 和 qmax 的值好确定，基本就是 8 个 bit 能表示的范围，这里取的分别是 -128 和 127（更详细的计算方式将会在下文的“静态量化”章节中描述）；Linear op 的权重为 torch.Tensor(\[\[ 0.4097, -0.2896, -0.4931\], \[-0.3738, -0.5541, 0.3243\]\])，因此其 min\_val 和 max\_val 分别为 -0.5541 和 0.4097，在这个上下文中，max\_val 将进一步取这俩绝对值的最大值。由此我们就可以得到：

*  scale = max\_val / (float(qmax - qmin) / 2) = 0.5541 / ((127 + 128) / 2) = 0.004345882...
 *  zp = 0

scale 和 zp 的计算细节还会在下文的“静态量化”章节中更详细的描述。从上面我们可以得知，权重部分的量化是“静态”的，是提前就转换完毕的，而之所以叫做“动态”量化，就在于前向推理的时候动态的把 input 的 float tensor 转换为量化 tensor。

在 forward的时候，nnqd.Linear 会调用 torch.ops.quantized.linear\_dynamic 函数，输入正是上面（pack 好后的）量化后的权重和 float 的 bias，而 torch.ops.quantized.linear\_dynamic 函数最终会被 PyTorch 分发到 C++ 中的 apply\_dynamic\_impl 函数，在这里，或者使用 FBGEMM 的实现（x86-64 设备），或者使用 QNNPACK 的实现（ARM 设备上）：

#ifdef USE_FBGEMM
    at::Tensor PackedLinearWeight::apply_dynamic_impl(at::Tensor input, bool reduce_range) {
      ...
      fbgemm::xxxx
      ...
    }
    #endif // USE_FBGEMM
    
    #ifdef USE_PYTORCH_QNNPACK
    at::Tensor PackedLinearWeightsQnnp::apply_dynamic_impl(at::Tensor input) {
      ...
      qnnpack::qnnpackLinearDynamic(xxxx)
      ...
    }
    #endif // USE_PYTORCH_QNNPACK

等等，input 还是 float32 的啊，这怎么运算嘛。别急，在上述的 apply\_dynamic\_impl 函数中，会使用下面的逻辑对输入进行量化：

Tensor q_input = at::quantize_per_tensor(input_contig, q_params.scale, q_params.zero_point, c10::kQUInt8);

也就是说，动态量化的本质就藏身于此：基于运行时对数据范围的观察，来动态确定对输入进行量化时的 scale 值。这就确保 input tensor 的 scale 因子能够基于输入数据进行优化，从而获得颗粒度更细的信息。

而模型的参数则是提前就转换为了 INT8 的格式（在使用 quantize\_dynamic API 的时候）。这样，当输入也被量化后，网络中的运算就使用向量化的 INT8 指令来完成。而在当前 layer 输出的时候，我们还需要把结果再重新转换为 float32——re-quantization 的 scale 值是依据 input、 weight 和 output scale 来确定的，定义如下：

**requant\_scale = input\_scale\_fp32 \* weight\_scale\_fp32 / output\_scale\_fp32**

实际上，在 apply\_dynamic\_impl 函数中，requant\_scales 就是这么实现的：

auto output_scale = 1.f
    auto inverse_output_scale = 1.f /output_scale;
    requant_scales[i] = (weight_scales_data[i] * input_scale) * inverse_output_scale;

这就是为什么在前面 Gemfield 提到过，经过量化版的 fc 的输出为torch.Tensor(\[\[\[\[-1.3038, -0.3847\], \[1.2856, 0.3969\]\]\]\])，已经变回正常的 float tensor 了。所以动态量化模型的前向推理过程可以概括如下：

#原始的模型，所有的tensor和计算都是浮点型
    previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                     /
    linear_weight_fp32
    
    #动态量化后的模型，Linear和LSTM的权重是int8
    previous_layer_fp32 -- linear_int8_w_fp32_inp -- activation_fp32 -- next_layer_fp32
                         /
       linear_weight_int8

总结下来，我们可以这么说：Post Training Dynamic Quantization，简称为 Dynamic Quantization，也就是动态量化，或者叫作Weight-only的量化，是提前把模型中某些 op 的参数量化为 INT8，然后在运行的时候动态的把输入量化为 INT8，然后在当前 op 输出的时候再把结果 requantization 回到 float32 类型。动态量化默认只适用于 Linear 以及 RNN 的变种。

##     ##

**04. Post Training Static Quantization**

与其介绍 post training static quantization 是什么，我们不如先来说明下它和 dynamic quantization 的相同点和区别是什么。相同点就是，都是把网络的权重参数转从 float32 转换为 int8；不同点是，需要把训练集或者和训练集分布类似的数据喂给模型（注意没有反向传播），然后通过每个 op 输入的分布特点来计算 activation 的量化参数（scale 和 zp）——称之为 Calibrate（定标）。

是的，静态量化包含有 activation 了，也就是 post process，也就是 op forward 之后的后处理。为什么静态量化需要 activation 呢？因为静态量化的前向推理过程自（始+1）至（终-1）都是 INT 计算，activation 需要确保一个 op 的输入符合下一个 op 的输入。

PyTorch 会使用五部曲来完成模型的静态量化：

**1. fuse\_model**

合并一些可以合并的 layer。这一步的目的是为了提高速度和准确度：

fuse_modules(model, modules_to_fuse, inplace=False, fuser_func=fuse_known_modules, fuse_custom_config_dict=None)

比如给 fuse\_modules 传递下面的参数就会合并网络中的 conv1、bn1、relu1：

torch.quantization.fuse_modules(gemfield_model, [['conv1', 'bn1', 'relu1']], inplace=True)

一旦合并成功，那么原始网络中的 conv1 就会被替换为新的合并后的 module（因为其是 list 中的第一个元素），而 bn1、relu1（list 中剩余的元素）会被替换为 nn.Identity()，这个模块是个占位符，直接输出输入。举个例子，对于下面的一个小网络：

class CivilNet(nn.Module):
        def __init__(self):
            super(CivilNet, self).__init__()
            syszuxin = 1
            syszuxout = 1
            self.conv = nn.Conv2d(syszuxin, syszuxout, kernel_size=1, stride=1, padding=0, groups=1, bias=False)
            self.fc = nn.Linear(3, 2,bias=False)
            self.relu = nn.ReLU(inplace=False)
    
        def forward(self, x):
            x = self.conv(x)
            x = self.fc(x)
            x = self.relu(x)
            return x

网络结构如下：

CivilNet(
      (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (fc): Linear(in_features=3, out_features=2, bias=False)
      (relu): ReLU()
    )

经过 torch.quantization.fuse\_modules(c, \[\[**'fc', 'relu'**\]\], inplace=**True**)后，网络变成了：

CivilNet(
      (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (fc): LinearReLU(
        (0): Linear(in_features=3, out_features=2, bias=False)
        (1): ReLU()
      )
      (relu): Identity()
    )

modules\_to\_fuse 参数的 list 可以包含多个 item list，或者是 submodule 的 op list 也可以，比如：\[ \['conv1', 'bn1', 'relu1'\], \['submodule.conv', 'submodule.relu'\]\]。有的人会说了，我要 fuse的module 被 Sequential 封装起来了，如何传参？参考下面的代码：

torch.quantization.fuse_modules(a_sequential_module, ['0', '1', '2'], inplace=True)

不是什么类型的 op 都可以参与合并，也不是什么样的顺序都可以参与合并。就目前来说，截止到 pytorch 1.7.1，只有如下的 op 和顺序才可以：

*  Convolution, Batch normalization
 *  Convolution, Batch normalization, Relu
 *  Convolution, Relu
 *  Linear, Relu
 *  Batch normalization, Relu

实际上，这个 mapping 关系就定义在 DEFAULT\_OP\_LIST\_TO\_FUSER\_METHOD 中：

DEFAULT_OP_LIST_TO_FUSER_METHOD : Dict[Tuple, Union[nn.Sequential, Callable]] = {
        (nn.Conv1d, nn.BatchNorm1d): fuse_conv_bn,
        (nn.Conv1d, nn.BatchNorm1d, nn.ReLU): fuse_conv_bn_relu,
        (nn.Conv2d, nn.BatchNorm2d): fuse_conv_bn,
        (nn.Conv2d, nn.BatchNorm2d, nn.ReLU): fuse_conv_bn_relu,
        (nn.Conv3d, nn.BatchNorm3d): fuse_conv_bn,
        (nn.Conv3d, nn.BatchNorm3d, nn.ReLU): fuse_conv_bn_relu,
        (nn.Conv1d, nn.ReLU): nni.ConvReLU1d,
        (nn.Conv2d, nn.ReLU): nni.ConvReLU2d,
        (nn.Conv3d, nn.ReLU): nni.ConvReLU3d,
        (nn.Linear, nn.ReLU): nni.LinearReLU,
        (nn.BatchNorm2d, nn.ReLU): nni.BNReLU2d,
        (nn.BatchNorm3d, nn.ReLU): nni.BNReLU3d,
    }

**2. 设置qconfig**

qconfig 是要设置到模型或者模型的子 module 上的。前文 Gemfield 就已经说过，qconfig 是 QConfig 的一个实例，QConfig 这个类就是维护了两个 observer，一个是 activation 所使用的 observer，一个是 op 权重所使用的 observer。

#如果要部署在x86 server上
    gemfield_model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    
    #如果要部署在ARM上
    gemfield_model.qconfig = torch.quantization.get_default_qconfig('qnnpack')

如果是 x86 和 arm 之外呢？抱歉，目前不支持。实际上，这里的 get\_default\_qconfig 函数的实现如下所示：

def get_default_qconfig(backend='fbgemm'):
        if backend == 'fbgemm':
            qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=True),weight=default_per_channel_weight_observer)
        elif backend == 'qnnpack':
            qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=False),weight=default_weight_observer)
        else:
            qconfig = default_qconfig
        return qconfig

default\_qconfig 实际上是 QConfig(activation=default\_observer, weight=default\_weight\_observer)，所以 gemfield 这里总结了一个表格：

<table> 
 <tbody> 
 <tr> 
 <td colspan="1">量化的backend </td> 
 <td colspan="1">activation</td> 
 <td colspan="1">weight</td> 
 </tr> 
 <tr> 
 <td colspan="1">fbgemm</td> 
 <td colspan="1">HistogramObserver (reduce_range=True)</td> 
 <td colspan="1">PerChannelMinMaxObserver&nbsp;(default_per_channel_weight_observer)</td> 
 </tr> 
 <tr> 
 <td colspan="1">qnnpack</td> 
 <td colspan="1">HistogramObserver (reduce_range=False)</td> 
 <td colspan="1">MinMaxObserver(default_weight_observer)</td> 
 </tr> 
 <tr> 
 <td colspan="1">默认（非fbgemm和qnnpack）</td> 
 <td colspan="1">MinMaxObserver(default_observer)</td> 
 <td colspan="1">MinMaxObserver(default_weight_observer)</td> 
 </tr> 
 </tbody> 
</table>

**3. prepare**

prepare 调用是通过如下 API 完成的：

gemfield_model_prepared = torch.quantization.prepare(gemfield_model)

prepare 用来给每个子 module 插入 Observer，用来收集和定标数据。以 activation 的 observer 为例，就是期望其观察输入数据得到四元组中的 min\_val 和 max\_val，至少观察个几百个迭代的数据吧，然后由这四元组得到 scale 和 zp 这两个参数的值。

module 上安插 activation 的 observer 是怎么实现的呢？还记得 \[1\] 一文中说过的“\_forward\_hooks 是通过 register\_forward\_hook 来完成注册的。这些 hooks 是在 forward 完之后被调用的......”吗？没错，CivilNet 模型中的 Conv2d、Linear、ReLU、QuantStub 这些 module 的 \_forward\_hooks 上都被插入了 activation 的 HistogramObserver，当这些子 module 计算完毕后，结果会被立刻送到其 \_forward\_hooks 中的 HistogramObserver 进行观察。

这一步完成后，CivilNet 网络就被改造成了：

CivilNet(
      (conv): Conv2d(
        1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False
        (activation_post_process): HistogramObserver()
      )
      (fc): Linear(
        in_features=3, out_features=2, bias=False
        (activation_post_process): HistogramObserver()
      )
      (relu): ReLU(
        (activation_post_process): HistogramObserver()
      )
      (quant): QuantStub(
        (activation_post_process): HistogramObserver()
      )
      (dequant): DeQuantStub()
    )

**4. 喂数据**

这一步不是训练。是为了获取数据的分布特点，来更好的计算 activation 的  scale 和 zp。至少要喂上几百个迭代的数据。

#至少观察个几百迭代
    for data in data_loader:
        gemfield_model_prepared(data)

**5. 转换模型**

第四步完成后，各个 op 权重的四元组（min\_val，max\_val，qmin, qmax）中的 min\_val，max\_val 已经有了，各个 op activation 的四元组（min\_val，max\_val，qmin, qmax）中的 min\_val，max\_val 也已经观察出来了。那么在这一步我们将调用 convert API：

gemfield_model_prepared_int8 = torch.quantization.convert(gemfield_model_prepared)

这个过程和 dynamic 量化类似，本质就是检索模型中 op 的 type，如果某个 op 的 type 属于字典 DEFAULT\_STATIC\_QUANT\_MODULE\_MAPPINGS 的 key（注意字典和动态量化的不一样了），那么，这个 op 将被替换为 key 对应的  value：

DEFAULT_STATIC_QUANT_MODULE_MAPPINGS = {
        QuantStub: nnq.Quantize,
        DeQuantStub: nnq.DeQuantize,
        nn.BatchNorm2d: nnq.BatchNorm2d,
        nn.BatchNorm3d: nnq.BatchNorm3d,
        nn.Conv1d: nnq.Conv1d,
        nn.Conv2d: nnq.Conv2d,
        nn.Conv3d: nnq.Conv3d,
        nn.ConvTranspose1d: nnq.ConvTranspose1d,
        nn.ConvTranspose2d: nnq.ConvTranspose2d,
        nn.ELU: nnq.ELU,
        nn.Embedding: nnq.Embedding,
        nn.EmbeddingBag: nnq.EmbeddingBag,
        nn.GroupNorm: nnq.GroupNorm,
        nn.Hardswish: nnq.Hardswish,
        nn.InstanceNorm1d: nnq.InstanceNorm1d,
        nn.InstanceNorm2d: nnq.InstanceNorm2d,
        nn.InstanceNorm3d: nnq.InstanceNorm3d,
        nn.LayerNorm: nnq.LayerNorm,
        nn.LeakyReLU: nnq.LeakyReLU,
        nn.Linear: nnq.Linear,
        nn.ReLU6: nnq.ReLU6,
        # Wrapper Modules:
        nnq.FloatFunctional: nnq.QFunctional,
        # Intrinsic modules:
        nni.BNReLU2d: nniq.BNReLU2d,
        nni.BNReLU3d: nniq.BNReLU3d,
        nni.ConvReLU1d: nniq.ConvReLU1d,
        nni.ConvReLU2d: nniq.ConvReLU2d,
        nni.ConvReLU3d: nniq.ConvReLU3d,
        nni.LinearReLU: nniq.LinearReLU,
        nniqat.ConvBn1d: nnq.Conv1d,
        nniqat.ConvBn2d: nnq.Conv2d,
        nniqat.ConvBnReLU1d: nniq.ConvReLU1d,
        nniqat.ConvBnReLU2d: nniq.ConvReLU2d,
        nniqat.ConvReLU2d: nniq.ConvReLU2d,
        nniqat.LinearReLU: nniq.LinearReLU,
        # QAT modules:
        nnqat.Linear: nnq.Linear,
        nnqat.Conv2d: nnq.Conv2d,
    }

替换的过程也和 dynamic 一样，使用 from\_float() API，这个 API 会使用前面的四元组信息计算出 op 权重和 op activation 的 scale 和 zp，然后用于量化。动态量化”章节时 Gemfield 说过要再详细介绍下 scale 和 zp 的计算过程，好了，就在这里。这个计算过程覆盖了如下的几个问题：

*  QuantStub 的 scale 和 zp 是怎么来的（静态量化需要插入 QuantStub，后文有说明）？
 *  conv activation 的 scale 和 zp 是怎么来的？
 *  conv weight 的 scale 和 zp 是怎么来的？
 *  fc activation 的 scale 和 zp 是怎么来的？
 *  fc weight 的 scale 和 zp 是怎么来的？
 *  relu activation 的 scale 和 zp 是怎么来的？
 *  relu weight 的...等等，relu 没有 weight。

我们就从 conv 来说起吧，还记得前面说过的 Observer 吗？分为 activation 和 weight 两种。以 Gemfield 这里使用的 fbgemm 后端为例，activation默认的observer 是 HistogramObserver、weight 默认的 observer 是 PerChannelMinMaxObserver。而计算 scale 和 zp 所需的四元组都是这些 observer 观察出来的呀（好吧，其中两个）。

在 convert API 调用中，pytorch 会将 Conv2d op 替换为对应的 QuantizedConv2d，在这个替换的过程中会计算 QuantizedConv2d activation 的 scale 和 zp 以及 QuantizedConv2d weight 的 scale 和 zp。

在各种 observer 中，计算 scale 和 zp 离不开这四个变量：min\_val，max\_val，qmin, qmax，分别代表输入的数据/权重的数据分布的最小值和最大值，以及量化后的取值范围的最小、最大值。qmin 和 qmax 的值好确定，基本就是 8 个 bit 能表示的范围，在pytorch中，qmin 和 qmax 是使用如下方式确定的：

if self.dtype == torch.qint8:
        if self.reduce_range:
            qmin, qmax = -64, 63
        else:
            qmin, qmax = -128, 127
    else:
        if self.reduce_range:
            qmin, qmax = 0, 127
        else:
            qmin, qmax = 0, 255

比如 conv 的 activation 的 observer（quint8）是 HistogramObserver，又是  reduce\_range 的，因此其 qmin,qmax = 0，127，而 conv 的 weight（qint8）是 PerChannelMinMaxObserver，不是 reduce\_range 的，因此其 qmin, qmax = -128, 127。

那么 min\_val，max\_val 又是怎么确定的呢？对于 HistogramObserver，其由输入数据 + 权重值根据 L2Norm(An approximation for L2 error minimization)确定；对于 PerChannelMinMaxObserver 来说，其由输入数据的最小值和最大值确定，比如在上述的例子中，值就是 -0.7898 和 -0.7898。

既然现在 conv weight 的 min\_val，max\_val，qmin, qmax 分别为 -0.7898、-0.7898、-128、 127，那如何得到 scale 和 zp 呢？PyTorch 就是用下面的逻辑进行计算的：

#qscheme 是 torch.per_tensor_symmetric 或者torch.per_channel_symmetric时
    max_val = torch.max(-min_val, max_val)
    scale = max_val / (float(qmax - qmin) / 2)
    scale = torch.max(scale, torch.tensor(self.eps, device=device, dtype=scale.dtype))
    if self.dtype == torch.quint8:
        zero_point = zero_point.new_full(zero_point.size(), 128)
    
    #qscheme 是 torch.per_tensor_affine时
    scale = (max_val - min_val) / float(qmax - qmin)
    scale = torch.max(scale, torch.tensor(self.eps, device=device, dtype=scale.dtype))
    zero_point = qmin - torch.round(min_val / scale)
    zero_point = torch.max(zero_point, torch.tensor(qmin, device=device, dtype=zero_point.dtype))
    zero_point = torch.min(zero_point, torch.tensor(qmax, device=device, dtype=zero_point.dtype))

由此 conv2d weight 的谜团就被我们解开了：

*  scale = 0.7898 / ((127 + 128)/2 ) = 0.0062
 *  zp = 0

再说说 QuantStub 的 scale 和 zp 是如何计算的。QuantStub 使用的是 HistogramObserver，根据输入从 \[-3,3\] 的分布，HistogramObserver 计算得到min\_val、max\_val 分别是 -3、2.9971，而 qmin 和 qmax 又分别是 0、127，其 schema 为 per\_tensor\_affine，因此套用上面的 per\_tensor\_affine 逻辑可得：

*  scale = (2.9971 + 3) / (127 - 0) = 0.0472
 *  zp = 0 - round(-3 /0.0472) = 64

其它计算同理，不再赘述。有了scale 和 zp，就有了量化版本的 module，上面那个 CivilNet 网络，经过静态量化后，网络的变化如下所示：

#原始的CivilNet网络：
    CivilNet(
      (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (fc): Linear(in_features=3, out_features=2, bias=False)
      (relu): ReLU()
    )
    
    #静态量化后的CivilNet网络：
    CivilNet(
      (conv): QuantizedConv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), scale=0.0077941399067640305, zero_point=0, bias=False)
      (fc): QuantizedLinear(in_features=3, out_features=2, scale=0.002811126410961151, zero_point=14, qscheme=torch.per_channel_affine)
      (relu): QuantizedReLU()
    )

**静态量化模型如何推理？**

我们知道，在 PyTorch 的网络中，前向推理逻辑都是实现在了每个 op 的 forward 函数中（参考：Gemfield：详解 Pytorch 中的网络构造 \[1\]）。而在 convert 完成后，所有的 op 被替换成了量化版本的 op，那么量化版本的 op 的 forward 会有什么不一样的呢？还记得吗？

**动态量化**中可是只量化了 op 的权重哦，输入的量化所需的 scale 的值是在推理过程中动态计算出来的。而**静态量化**中，统统都是提前就计算好的。我们来看一个典型的静态量化模型的推理过程：

import torch
    import torch.nn as nn
    
    class CivilNet(nn.Module):
        def __init__(self):
            super(CivilNet, self).__init__()
            in_planes = 1
            out_planes = 1
            self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, padding=0, groups=1, bias=False)
            self.fc = nn.Linear(3, 2,bias=False)
            self.relu = nn.ReLU(inplace=False)
            self.quant = QuantStub()
            self.dequant = DeQuantStub()
    
        def forward(self, x):
            x = self.quant(x)
            x = self.conv(x)
            x = self.fc(x)
            x = self.relu(x)
            x = self.dequant(x)
            return x

网络 forward 的开始和结束还必须安插 QuantStub 和 DeQuantStub，如上所示。否则运行时会报错：RuntimeError: Could not run 'quantized::conv2d.new' with arguments from the 'CPU' backend. 'quantized::conv2d.new' is only available for these backends: \[QuantizedCPU\]。

QuantStub 在 observer 阶段会记录参数值，DeQuantStub 在 prepare阶段相当于 Identity；而在 convert API 调用过程中，会分别被替换为 nnq.Quantize 和 nnq.DeQuantize。在这个章节要介绍的推理过程中，QuantStub，也就是 nnq.Quantize 在做什么工作呢？如下所示：

def forward(self, X):
        return torch.quantize_per_tensor(X, float(self.scale), int(self.zero_point), self.dtype)

是不是呼应了前文中的“tensor 的量化”章节？这里的 scale 和 zero\_point 的计算方式前文也刚介绍过。而 nnq.DeQuantize 做了什么呢？很简单，把量化 tensor 反量化回来。

def forward(self, Xq):
        return Xq.dequantize()

是不是又呼应了前文中的“tensor的量化”章节？我们就以上面的 CivilNet 网络为例，当在静态量化后的模型进行前向推理和原始的模型的区别是什么呢？假设网络的输入为 torch.Tensor(\[\[\[\[-1,-2,-3\],\[1,2,3\]\]\]\])：

c = CivilNet()
    t = torch.Tensor([[[[-1,-2,-3],[1,2,3]]]])
    c(t)

假设 conv 的权重为 torch.Tensor(\[\[\[\[-0.7867\]\]\]\])，假设 fc 的权重为 torch.Tensor(\[\[ 0.4097, -0.2896, -0.4931\], \[-0.3738, -0.5541, 0.3243\]\])，那么在原始的 CivilNet 前向中，从输入到输出的过程依次为：

而在静态量化的模型前向中，总体情况如下：

#input
    torch.Tensor([[[[-1,-2,-3],[1,2,3]]]])
    
    #QuantStub后 (scale=tensor([0.0472]), zero_point=tensor([64]))
    tensor([[[[-0.9916, -1.9833, -3.0221],[ 0.9916,  1.9833,  3.0221]]]],
           dtype=torch.quint8, scale=0.04722102731466293, zero_point=64)
    
    #经过卷积后（权重为torch.Tensor([[[[-0.7898]]]], dtype=torch.qint8, scale=0.0062, zero_point=0))
    #conv activation（输入）的scale为0.03714831545948982，zp为64
    torch.Tensor([[[[ 0.7801,  1.5602,  2.3775],[-0.7801, -1.5602, -2.3775]]]], scale=0.03714831545948982, zero_point=64)
    
    #经过fc后（权重为torch.Tensor([[ 0.4100, -0.2901, -0.4951],[-0.3737, -0.5562,  0.3259]], dtype=torch.qint8, scale=tensor([0.0039, 0.0043]),zero_point=tensor([0, 0])) )
    #fc activation（输入）的scale为0.020418135449290276, zp为64
    torch.Tensor([[[[-1.3068, -0.3879],[ 1.3068,  0.3879]]]], dtype=torch.quint8, scale=0.020418135449290276, zero_point=64)
    
    #经过relu后
    torch.Tensor([[[[0.0000, 0.0000],[1.3068, 0.3879]]]], dtype=torch.quint8, scale=0.020418135449290276, zero_point=64)
    
    #经过DeQuantStub后
    torch.Tensor([[[[0.0000, 0.0000],[1.3068, 0.3879]]]])

Gemfield 这里用原始的 python 语句来分步骤来介绍下。首先是 **QuantStub** 的工作：

import torch
    import torch.nn.quantized as nnq
    #输入
    >>> x = torch.Tensor([[[[-1,-2,-3],[1,2,3]]]])
    >>> x
    tensor([[[[-1., -2., -3.],
              [ 1.,  2.,  3.]]]])
    
    #经过QuantStub
    >>> xq = torch.quantize_per_tensor(x, scale = 0.0472, zero_point = 64, dtype=torch.quint8)
    >>> xq
    tensor([[[[-0.9912, -1.9824, -3.0208],
              [ 0.9912,  1.9824,  3.0208]]]], size=(1, 1, 2, 3),
           dtype=torch.quint8, quantization_scheme=torch.per_tensor_affine,
           scale=0.0472, zero_point=64)
    
    >>> xq.int_repr()
    tensor([[[[ 43,  22,   0],
              [ 85, 106, 128]]]], dtype=torch.uint8)

我们特意在网络前面安插的 QuantStub 完成了自己的使命，其 scale = 0.0472、zero\_point = 64 是静态量化完毕后就已经知道的，然后通过 quantize\_per\_tensor 调用把输入的 float tensor 转换为了量化 tensor，然后送给接下来的 **Conv2d——量化版本的 Conv2d：**

>>> c = nnq.Conv2d(1,1,1)
    >>> weight = torch.Tensor([[[[-0.7898]]]])
    >>> qweight = torch.quantize_per_channel(weight, scales=torch.Tensor([0.0062]).to(torch.double), zero_points = torch.Tensor([0]).to(torch.int64), axis=0, dtype=torch.qint8)
    >>> c.set_weight_bias(qweight, None)
    >>> c.scale = 0.03714831545948982
    >>> c.zero_point = 64
    >>> x = c(xq)
    >>> x
    tensor([[[[ 0.7801,  1.5602,  2.3775],
              [-0.7801, -1.5602, -2.3775]]]], size=(1, 1, 2, 3),
           dtype=torch.quint8, quantization_scheme=torch.per_tensor_affine,
           scale=0.03714831545948982, zero_point=64)

同理，Conv2d 的权重的 scale=0.0062、zero\_points=0 是静态量化完毕就已知的，其 activation 的 scale = 0.03714831545948982、zero\_point = 64 也是量化完毕已知的。然后送给 nnq.Conv2d 的 forward 函数（参考：\[1\]），其 forward 逻辑为：

def forward(self, input):
        return ops.quantized.conv2d(input, self._packed_params, self.scale, self.zero_point)

Conv2d 计算完了，我们停下来反省一下。如果是按照浮点数计算，那么 -0.7898 \* -0.9912 大约是 0.7828，但这里使用 int8 的计算方式得到的值是 0.7801，这说明已经在引入误差了（大约为 0.34% 的误差）。这也是前面 gemfield 说的使用 fuse\_modules 可以提高精度的原因，因为每一层都会引入类似的误差。

后面 Linear 的计算同理，其 forward 逻辑为：

def forward(self, x):
        return torch.ops.quantized.linear(x, self._packed_params._packed_params, self.scale, self.zero_point)

可以看到，所有以量化方式计算完的值现在需要经过 activation 的计算。这是静态量化和动态量化的本质区别之一：op 和 op 之间不再需要转换回到 float tensor 了。通过上面的分析，我们可以把静态量化模型的前向推理过程概括为如下的形式：

#原始的模型，所有的tensor和计算都是浮点型
    previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                        /
        linear_weight_fp32
    
    #静态量化的模型，权重和输入都是int8
    previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                        /
      linear_weight_int8

最后再来描述下动态量化和静态量化的最大区别：

*  静态量化的 float 输入必经 QuantStub 变为 int，此后到输出之前都是 int；
 *  动态量化的 float 输入是经动态计算的 scale 和 zp 量化为 int，op 输出时转换回 float。

##     ##

**05. QAT（Quantization Aware Training）**

前面两种量化方法都有一个 post 关键字，意思是模型训练完毕后所做的量化。而 QAT 则不一样，是指在训练过程中就开启了量化功能。

QAT 需要五部曲，说到这里，你可能想到了静态量化，那不妨对比着来看。

**1. 设置qconfig**

在设置 qconfig 之前，模型首先设置为训练模式，这很容易理解，因为 QAT 的着力点就是 T 嘛：

cnet = CivilNet()
    cnet.train()

使用 get\_default\_qat\_qconfig API 来给要 QAT 的网络设置 qconfig：

cnet.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

不过，这个 qconfig 和静态量化中的可不一样啊。前文说过 qconfig 维护了两个observer，activation 的和权重的。QAT 的 qconfig 中，activation 和权重的observer 都变成了 FakeQuantize（和 observer 是 has a 的关系，也即包含一个 observer），并且参数不一样（qmin、qmax、schema,dtype,qschema,reduce\_range 这些参数），如下所示：

#activation的observer的参数
    FakeQuantize.with_args(observer=MovingAverageMinMaxObserver,quant_min=0,quant_max=255,reduce_range=True)
    
    #权重的observer的参数
    FakeQuantize.with_args(observer=MovingAveragePerChannelMinMaxObserver,
                                                                   quant_min=-128,
                                                                   quant_max=127,
                                                                   dtype=torch.qint8,
                                                                   qscheme=torch.per_channel_symmetric,
                                                                   reduce_range=False,
                                                                   ch_axis=0)

这里 FakeQuantize 包含的 observer 是 MovingAverageMinMaxObserver，继承自前面提到过的 MinMaxObserver，但是求最小值和最大值的方法有点区别，使用的是如下公式：

![4abb834a8734bbbcc57b74841946c981.png][]

*  Xmin、Xmax 是当前运行中正在求解和最终求解的最小值、最大值；
 *  X 是当前输入的 tensor；
 *  c 是一个常数，PyTorch 中默认为 0.01，也就是最新一次的极值由上一次贡献 99%，当前的 tensor 贡献 1%。

MovingAverageMinMaxObserver 在求 min、max 的方式和其基类 MinMaxObserver 有所区别之外，scale 和 zero\_points 的计算则是一致的。那么在包含了上述的 observer 之后，FakeQuantize 的作用又是什么呢？看下面的步骤。

**2. fuse\_modules**

和静态量化一样，不再赘述。

**3. prepare\_qat**

在静态量化中，我们这一步使用的是 prepare API，而在 QAT 这里使用的是 prepare\_qat API。最重要的区别有两点：

*  prepare\_qa t 要把 qconfig 安插到每个 op 上，qconfig 的内容本身就不同，参考五部曲中的第一步；
 *  prepare\_qat 中需要多做一步转换子 module 的工作，需要 inplace 的把模型中的一些子 module 替换了，替换的逻辑就是从 DEFAULT\_QAT\_MODULE\_MAPPINGS 的 key 替换为 value，这个字典的定义如下：

# Default map for swapping float module to qat modules
    DEFAULT_QAT_MODULE_MAPPINGS : Dict[Callable, Any] = {
        nn.Conv2d: nnqat.Conv2d,
        nn.Linear: nnqat.Linear,
        # Intrinsic modules:
        nni.ConvBn1d: nniqat.ConvBn1d,
        nni.ConvBn2d: nniqat.ConvBn2d,
        nni.ConvBnReLU1d: nniqat.ConvBnReLU1d,
        nni.ConvBnReLU2d: nniqat.ConvBnReLU2d,
        nni.ConvReLU2d: nniqat.ConvReLU2d,
        nni.LinearReLU: nniqat.LinearReLU
    }

因此，同静态量化的 prepare 相比，prepare\_qat 在多插入 fake\_quants、又替换了 nn.Conv2d、nn.Linear 之后，CivilNet 网络就被改成了如下的样子：

CivilNet(
      (conv): QATConv2d(
        1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False
        (activation_post_process): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
        (weight_fake_quant): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
      )
      (fc): QATLinear(
        in_features=3, out_features=2, bias=False
        (activation_post_process): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
        (weight_fake_quant): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
      )
      (relu): ReLU(
        (activation_post_process): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
      )
      (quant): QuantStub(
        (activation_post_process): FakeQuantize(
          fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
          (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
        )
      )
      (dequant): DeQuantStub()
    )

**4. 喂数据**

和静态量化完全不同，在 QAT 中这一步是用来训练的。我们知道，在 PyTorch 的网络中，前向推理逻辑都是实现在了每个op的 forward 函数中（参考：Gemfield：详解 Pytorch 中的网络构造 \[1\]）。而在 prepare\_qat 中，所有的 op 被替换成了 QAT 版本的 op，那么这些 op 的 forward 函数有什么特别的地方呢？

Conv2d 被替换为了 QATConv2d：

def forward(self, input):
       return self.activation_post_process(self._conv_forward(input, self.weight_fake_quant(self.weight)))

Linear 被替换为了 QATLinear:

def forward(self, input):
        return self.activation_post_process(F.linear(input, self.weight_fake_quant(self.weight), self.bias))

ReLU 还是那个 ReLU，不说了。总之，你可以看出来，每个 op 的输入都需要经过 self.weight\_fake\_quant 来处理下，输出又都需要经过 self.activation\_post\_process 来处理下，这两个都是 FakeQuantize 的实例，只是里面包含的 observer 不一样。以 Conv2d 为例：

#conv2d
 weight=functools.partial(<class 'torch.quantization.fake_quantize.FakeQuantize'>, 
            observer=<class 'torch.quantization.observer.MovingAveragePerChannelMinMaxObserver'>, 
            quant_min=-128, quant_max=127, dtype=torch.qint8, 
            qscheme=torch.per_channel_symmetric, reduce_range=False, ch_axis=0))
 
 activation=functools.partial(<class 'torch.quantization.fake_quantize.FakeQuantize'>, 
             observer=<class 'torch.quantization.observer.MovingAverageMinMaxObserver'>, 
             quant_min=0, quant_max=255, reduce_range=True)

而 FakeQuantize 的 forward 函数如下所示：

def forward(self, X):
            if self.observer_enabled[0] == 1:
                #使用移动平均算法计算scale和zp
    
            if self.fake_quant_enabled[0] == 1:
                X = torch.fake_quantize_per_channel_or_tensor_affine(X...)
            return X

FakeQuantize 中的 fake\_quantize\_per\_channel\_or\_tensor\_affine 实现了 quantize 和 dequantize，用公式表示的话为：out = (clamp(round(x/scale + zero\_point), quant\_min, quant\_max)-zero\_point)\*scale。也就是说，这是把量化的误差引入到了训练 loss 之中呀！

这样，在 QAT 中，所有的 weights 和 activations 就像上面那样被 fake quantized了，且参与模型训练中的前向和反向计算。float 值被 round 成了（用来模拟的）int8 值，但是所有的计算仍然是通过 float 来完成的。这样以来，所有的权重在优化过程中都能感知到量化带来的影响，称之为量化感知训练（支持 cpu 和 cuda），精度也因此更高。

**5. 转换**

这一步和静态量化一样，不再赘述。需要注意的是，QAT 中，有一些 module 在 prepare 中已经转换成新的 module 了，所以静态量化中所使用的字典包含有如下的条目：

DEFAULT_STATIC_QUANT_MODULE_MAPPINGS = {
        ......
        # QAT modules:
        nnqat.Linear: nnq.Linear,
        nnqat.Conv2d: nnq.Conv2d,
    }

总结下来就是：

# 原始的模型，所有的tensor和计算都是浮点
    previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                          /
        linear_weight_fp32
    
    # 训练过程中，fake_quants发挥作用
    previous_layer_fp32 -- fq -- linear_fp32 -- activation_fp32 -- fq -- next_layer_fp32
                               /
       linear_weight_fp32 -- fq
    
    # 量化后的模型进行推理，权重和输入都是int8
    previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                         /
       linear_weight_int8

**06. 总结**

那么如何更方便的在你的代码中使用 PyTorch 的量化功能呢？一个比较优雅的方式就是使用 deepvac 规范——这是一个定义了 PyTorch 工程标准的项目：

https://github.com/DeepVAC/deepvac

基于 deepvac 规范（包含库），我们只需要简单的打开几个开关就可以使用上述的三种量化功能。

***注：文仅学术分享，侵删***

**猜您喜欢：**

**[超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

[**附下载 | 《Python进阶》中文版**][_ _Python]

[**附下载 | 经典《Think Python》中文版**][_ _Think Python]

[**附下载 | 《Pytorch模型训练实用教程》**][_ _Pytorch]

[**附下载 | 最新2020李沐《动手学深度学习》**][_ _2020]

**[附下载 | 《可解释的机器学习》中文版][Link 1]**

[**附下载 |《TensorFlow 2.0 深度学习算法实战》**][_TensorFlow 2.0]

**[附下载 | 超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

**[附下载 |《计算机视觉中的数学方法》分享][Link 2]**

![0de1519735903edf23c6e4806bd0153b.png][]

[4abb834a8734bbbcc57b74841946c981.png]: /images/20221022/81c75f26aaeb41e49e92a24dcdf2f5e9.png
[100_CVPR 2020_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a94cc95e205ae880cb9518cb16a54c0abde69e5a8654f7fedf70a0cc990775265bf28a1c&idx=1&mid=2247486455&scene=21&sn=94156f53845cdf07e859e343862be745#wechat_redirect
[_ _Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29abb6c95e22a04d1eb019f964c8ae97b2f3af98c22f730f5d5ed2aea9fdb320b00c688967&idx=2&mid=2247485709&scene=21&sn=55a51c98ab48f393ba305c082ba0e4fe#wechat_redirect
[_ _Think Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a418c95e2d0e9a3de46e0e4d1b8d4947154d8639418a511a3091300aa31699932046556e&idx=2&mid=2247485091&scene=21&sn=f569d5d70c2e6c9cc567fbc5147787a4#wechat_redirect
[_ _Pytorch]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa78c95e236e380c9fd2ae79917d7451ccdd4803c5e6047a16f921d1b61d18efc79ada7c&idx=2&mid=2247485635&scene=21&sn=90608bce32fe353cf3c64b39e3406bc3#wechat_redirect
[_ _2020]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4a4c95e2db2aab16bdcefa78795bbb7efbc70b533bb197d73f617558f2f8a523120c6e3&idx=2&mid=2247484959&scene=21&sn=563999cc8d48498beed00270844ceec4#wechat_redirect
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4bac95e2dac758c4aac10a3aac49eb062bc0829d4ef2ebc8e6d5069d246385c5084a2b4&idx=1&mid=2247484929&scene=21&sn=523220858287ed3e6105919703c8e833#wechat_redirect
[_TensorFlow 2.0]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a5b0c95e2ca6599d4a91d5d40900566b7525de3dd85259e7a6c173aa26aa7ea3641d1aaf&idx=2&mid=2247485195&scene=21&sn=e3387067db29d1759755c6c2ea96c870#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e23993e057fb0dbe09f2810ea07184a439c173f404a7ab657cadcde3d24efa52a&idx=2&mid=2247485492&scene=21&sn=c4fb135de02cd52f06bc61b0e79476e4#wechat_redirect
[0de1519735903edf23c6e4806bd0153b.png]: /images/20221022/01858f961a7f4d0eaaa2db52b427e8a9.png