详解pytorch动态量化

前言

想要读取动态量化后模型的int8分布，但是发现模型内部已经是float，很诧异。。

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 ：

• 添加了 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的量化

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
<output>
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.qint8)
<output>
tensor([[0.5000, 0.5000, 0.5000],
        [1.0000, 0.0000, 0.5000]], size=(2, 3),dtype=torch.qint8,quantization_scheme=torch.per_tensor_affine, scale=0.5, zero_point=8)
>>> xq.int_repr()
<output>
tensor([[ 9,  9,  9],
        [10,  8,  9]], dtype=torch.int8)

quantize_per_tensor 函数就是使用给定的 scale 和 zero_point来把一个 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 qint8
>>> 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)\ast scale$

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

• 量化会有精度损失；
• 选择合适的 scale 和 zp 可以有效降低精度损失。

可以把 scale 和 zp 分别换成 scale = 0.0036, zero_point = 0试试

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

Tensor 的量化支持两种模式：

• 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就是网络在前向推理的时候动态的量化 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（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,
            }

这就是为什么动态量化只量化 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

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

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。

那么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

从上面我们可以得知，权重部分的量化是“静态”的，是提前就转换完毕的，而之所以叫做“动态”量化，就在于前向推理的时候动态的把 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 来确定的，定义如下：$requan{t}_{s}cale=input\_scale\_fp32\ast 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 的变种。

参考资料

https://blog.csdn.net/c9Yv2cf9I06K2A9E/article/details/113488003