如何使用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 的步骤如下：

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

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)
  (