pytorch中模型量化的prepare和convert API

1. prepare函数

使用prepare函数可以将模型准备好以进行后续的量化操作。为每个module插入Observer，用来收集和定标数据。prepare好的模型可以通过使用torch.quantization.convert函数将模型转换为量化模型。

def prepare(model, inplace=False, allow_list=None,
            observer_non_leaf_module_list=None,
            prepare_custom_config_dict=None):

函数的返回值取决于inplace参数的设置。如果inplace为True，则该函数不返回任何值，直接在原始模型上进行操作；如果inplace为False，则会返回一个新的量化模型。

参数：

• model：输入的需要量化的模型
• inplace（可选）：一个布尔值，指定是否在原始模型上进行操作。如果设置为True，则会就地修改原始模型；如果设置为False（默认值），则会返回一个新的量化模型。
• allow_list（可选）：一个可迭代对象，包含需要参与量化的模块类型。如果设置为None（默认值），则会自动选择适合量化的模块类型。如果提供了一个类型列表，则只有这些类型的模块会被量化。
• observer_non_leaf_module_list（可选）：一个可迭代对象，包含需要在非叶子模块上插入观察器的模块类型。默认情况下，只有叶子模块上的观察器会被插入。如果想要在非叶子模块上使用观察器，则需要提供相应的模块类型列表。
• prepare_custom_config_dict（可选）：一个字典，用于提供自定义的准备配置。可以用于配置特定模块类型的量化行为。

module上安插入activation的observer的实现可以参考Gemfield​：PyTorch的量化。
模型中的Conv2d、Linear、ReLU、QuantStub这些module的_forward_hooks上都被插入了activation的HistogramObserver，当这些子module计算完毕后，结果会被立刻送到其_forward_hooks中的HistogramObserver进行观察。

prepare量化模型仍然是浮点数类型。要将量化模型转换为量化权重和激活函数，可以使用torch.quantization.convert函数。

2. convert函数

将模型转换为量化模型。它会应用预定义的或自定义的量化方法，并将模型中的权重和激活函数转换为量化表示，从而减少计算和内存开销。转换后的量化模型可以在量化推理中使用，以获得更高的效率和性能。

def convert(
        module, mapping=None, inplace=False, remove_qconfig=True,
        is_reference=False, convert_custom_config_dict=None):

参数：

• module：要转换为量化模型的模块，通常是一个PyTorch模型的实例。
• mapping（可选）：一个字典，用于指定特定模块类型的量化方法。字典的键是模块类型，值是用于量化的量化方法。如果未提供映射，将使用默认的量化方法。
• inplace（可选）：一个布尔值，指定是否在原始模型上进行操作。如果设置为True，则会就地修改原始模型；如果设置为False（默认值），则会返回一个新的量化模型。
• remove_qconfig（可选）：一个布尔值，指定是否在转换后删除量化配置信息。如果设置为True（默认值），将删除模型中的量化配置信息；如果设置为False，将保留量化配置信息。
• is_reference（可选）：一个布尔值，指定是否将模型转换为参考模式。参考模式是一种特殊的量化模式，用于生成具有高精度参考的量化模型。默认为False，表示转换为正常的量化模型。
• convert_custom_config_dict（可选）：一个字典，用于提供自定义的转换配置。可以用于配置特定模块类型的量化行为

本质就是检索模型中op的type，如果某个op的type属于字典mapping（默认的在torch/ao/quantization/quantization_mappings.py中），这个op就会被替换将key对应的value，如nn.Conv2d被替换为 nnq.Conv2d。

并且之前Observer观测的参数将在这一步进行计算获得scale和zp

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

3. QuantStub和DeQuantStub

参考了Gemfield​：PyTorch的量化。
在量化流程中，通常需要在模型的输入和输出之间插入量化和反量化操作。

3.1 QuantStub

用于在模型的输入处插入量化操作，将浮点数输入转换为量化表示。量化表示是一种离散的表示形式，用于减少计算和内存开销。

在校准（calibration）前和observer一样
在convert API调用过程中，会被替换为nnq.Quantize

class QuantStub(nn.Module):
    def __init__(self, qconfig=None):
        super(QuantStub, self).__init__()
        if qconfig:
            self.qconfig = qconfig

    def forward(self, x):
        return x

• 参数qconfig用于设置量化配置。如果提供了qconfig，则将其赋值给self.qconfig属性。没有提供则使用默认的qconfig。
• QuantStub类实现中的forward方法只是简单地返回输入数据，没有执行任何实际的量化操作。这样的实现可以将QuantStub用作占位符，在量化模型中标记需要进行量化的位置。

QuantStub在前向推理中相当于通过quantize_per_tensor调用把输入的float tensor转换为了量化tensor，然后送给接下来的op（如Conv2d–量化版本的Conv2d）

其中QuantStub的s=0.0472和zp=64是是静态量化完毕后就已知的，然后通过quantize_per_tensor调用把输入的float tensor转换为了量化tensor，再传递给下一层。

在使用quantize_per_tensor时会引入一定的量化误差，

In[7]: c = nnq.Conv2d(1,1,1)
In[8]: weight = torch.Tensor([[[[-0.7898]]]])
In[9]: qweight = torch.quantize_per_channel(weight, In[10]:scales=torch.Tensor([0.0062]).to(torch.double), zero_points = torch.Tensor([0]).to(torch.int64), axis=0, dtype=torch.qint8)
In[10]: c.set_weight_bias(qweight, None)
In[11]: c.scale = 0.03714831545948982
In[12]: c.zero_point = 64
In[13]: x = c(xq)
In[14]: x
Out[14]: 
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函数

3.2 DeQuantStub

用于在模型的输出处插入反量化操作，将量化表示转换回浮点数表示。

在校准（calibration）前和identity一样
在convert API调用过程中，会被替换为nnq.DeQuantize。

class DeQuantStub(nn.Module):
    def __init__(self, qconfig=None):
        super(DeQuantStub, self).__init__()
        if qconfig:
            self.qconfig = qconfig

    def forward(self, x):
        return x