如何使用PyTorch的量化功能？

背景

在深度学习中，量化指的是使用更少的 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 的量化。

Tensor的量化

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

代码语言：javascript复制

>>> 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 的关系大概就是:

代码语言：javascript复制

xq = round(x / scale   zero_point)

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

代码语言：javascript复制

xq = round(zero_point) = zero_point

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

代码语言：javascript复制

# 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。也就是：

代码语言：javascript复制

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 量化。

Post Training Dynamic Quantization

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

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

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

代码语言：javascript复制

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 就会使用如下的默认值：

代码语言：javascript复制

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 的一个实例，使用如下的参数进行构造：

代码语言：javascript复制

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 使用一个小网络来演示下：

代码语言：javascript复制

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

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

代码语言：javascript复制

#原始网络
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：

代码语言：javascript复制

# 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() 方法中，通过如下方式实例化：

代码语言：javascript复制

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 的打包，在源代码中是这样的：

代码语言：javascript复制

#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。量化完后的模型在推理的时候有什么不一样的呢？在原始网络中，从输入到最终输出是这么计算的：

代码语言：javascript复制

#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]]]])

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

代码语言：javascript复制

#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 设备上）：

代码语言：javascript复制

#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 函数中，会使用下面的逻辑对输入进行量化：

代码语言：javascript复制

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 就是这么实现的：

代码语言：javascript复制

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 了。所以动态量化模型的前向推理过程可以概括如下：

代码语言：javascript复制

#原始的模型，所有的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 的变种。

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。这一步的目的是为了提高速度和准确度：

代码语言：javascript复制

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

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

代码语言：javascript复制

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

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

代码语言：javascript复制

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

网络结构如下：

代码语言：javascript复制

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)后，网络变成了：

代码语言：javascript复制

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 封装起来了，如何传参？参考下面的代码：

代码语言：javascript复制

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 中：

代码语言：javascript复制

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。

代码语言：javascript复制

#如果要部署在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 函数的实现如下所示：

代码语言：javascript复制

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 这里总结了一个表格：

量化的backend	activation	weight
fbgemm	HistogramObserver(reduce_range=True)	PerChannelMinMaxObserver (default_per_channel_weight_observer)
qnnpack	HistogramObserver(reduce_range=False)	MinMaxObserver(default_weight_observer)
默认（非fbgemm和qnnpack）	MinMaxObserver(default_observer)	MinMaxObserver(default_weight_observer)

3. prepare

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

代码语言：javascript复制

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 网络就被改造成了：

代码语言：javascript复制

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。至少要喂上几百个迭代的数据。

代码语言：javascript复制

#至少观察个几百迭代
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：

代码语言：javascript复制

gemfield_model_prepared_int8 = torch.quantization.convert(gemfield_model_prepared)

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

代码语言：javascript复制

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 是使用如下方式确定的：

代码语言：javascript复制

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 就是用下面的逻辑进行计算的：

代码语言：javascript复制

#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 网络，经过静态量化后，网络的变化如下所示：

代码语言：javascript复制

#原始的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 的值是在推理过程中动态计算出来的。而静态量化中，统统都是提前就计算好的。我们来看一个典型的静态量化模型的推理过程：

代码语言：javascript复制

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 在做什么工作呢？如下所示：

代码语言：javascript复制

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 反量化回来。

代码语言：javascript复制

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

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

代码语言：javascript复制

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 前向中，从输入到输出的过程依次为：

代码语言：javascript复制

#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]]]])

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

代码语言：javascript复制

#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 的工作：

代码语言：javascript复制

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：

代码语言：javascript复制

>>> 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 逻辑为：

代码语言：javascript复制

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 逻辑为：

代码语言：javascript复制

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 了。通过上面的分析，我们可以把静态量化模型的前向推理过程概括为如下的形式：

代码语言：javascript复制

#原始的模型，所有的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。

QAT（Quantization Aware Training）

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

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

1. 设置qconfig

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

代码语言：javascript复制

cnet = CivilNet()
cnet.train()

使用 get_default_qat_qconfig API 来给要 QAT 的网络设置 qconfig：

代码语言：javascript复制

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 这些参数），如下所示：

代码语言：javascript复制

#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，但是求最小值和最大值的方法有点区别，使用的是如下公式：

• 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，这个字典的定义如下：

代码语言：javascript复制

# 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 网络就被改成了如下的样子：

代码语言：javascript复制

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：

代码语言：javascript复制

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

Linear 被替换为了 QATLinear:

代码语言：javascript复制

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 为例：

代码语言：javascript复制

#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 函数如下所示：

代码语言：javascript复制

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 了，所以静态量化中所使用的字典包含有如下的条目：

代码语言：javascript复制

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

总结下来就是：

代码语言：javascript复制

# 原始的模型，所有的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

总结

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

https://github.com/DeepVAC/deepvac

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

参考文献

[1] https://zhuanlan.zhihu.com/p/53927068

代码语言：javascript复制

git 编程算法 pytorch api

0 人点赞