开个新坑速通一下量化，目标最终实现完善的DCCRN量化

由于DCCRN实现较为复杂，我们先从简单的ResNet入手

量化对象与原理

在对一个模型进行量化前，我们肯定首先需要理清模型本身的结构、流程和涉及组件，并分析哪里可能是模型计算的瓶颈区。对于任何一个模型，我们的量化对象主要可以分为：

网络组件：即Linear，卷积、LSTM或是attention这类网络的核心构件，也是主要的量化对象
norm：即各类标准化，神经网络中所有标准化（BN，LN，GN，IN）表达公式都可以写为：
$y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$
其中γ和β是可学习参数，量化norm的意义就是将norm计算fold进前面的网络权重中，这样理论上可以直接省去BN这个计算（包括减均值、除标准差等，后面细说），减少算子数量、内存独写和推理延迟
激活值：将激活值的范围映射到当前的量化分布上
池化：也是将输出数值控制在当前的量化范围内，比较容易实现
其他：由网络本身结构确定，例如残差连接，跳跃连接，也需要控制数值范围

核心公式

对称量化

$q=\text{round}(\frac x s)$

其中round为四舍五入，s为量化因子scale：

$s=\frac {max(|x_{min}|,|x_{max}|)}{127}$

由于分子是取最大值，因此量化范围为对称范围；例如原本x∈[-0.8,1]，对称量化的实际范围仍为[-1,1]，会浪费一部分区间

非对称量化

$q=\text{round}(\frac x s)+z$ $s=\frac {x_{max}-x_{min}}{255}$ $z=\text{round}(\frac{q_{min}-x_{min}}{s})$

其中q_min通常就是0

非对称量化可以更充分利用原本的数据区间，量化范围就等于数据本身的最小值到最大值

使用对象

对于神经网络的权重，其本身分布就接近对称，因此可以直接使用简单的对称量化（当然具体要看情况）；而对于激活值，很多激活值本身大部分范围（或全部）都是≥0的，因此要使用非对称

反量化

$x≈s(q-z)$

per channel和per tensor计算

就是scale的计算范围是逐通道还是直接在整个张量上算，对于卷积权重，每个通道可能分布差异较大，优先per channel；对于激活值，则优先效率更高的per tensor

BN folding

对于Batch norm，其通常是配合卷积一起出现的，由norm的公式，我们可以整理为：

$y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta=ax_{convOut}+b$

而卷积本身又是：

$x_{convOut}=W *x+b_{conv}$

那么BN本身其实就可以融合进卷积当中：

$y=aW*x+(ab_{conv}+b)=W'*x+b'$

关于BN和其他norm的不同

在默认情况下，BN在推理阶段使用的方差和均值均是在训练过程中计算的累计均值，是固定值，因此在推理阶段其本身就是一个线性变换；而其他的norm在推理过程中使用的均值方差由当前输入值计算而来，folding难度更高

模型评估

在量化之前，先评估模型每一层的参数量，计算量，明确我们的量化对象，对于参数量大，计算量大的模块是我们的主要量化目标

量化后端

目前我的PC只有FBGEMM和ONEDNN

完全调用torch API实现PTQ

在该节中，我们使用torch 的FX Graph Mode相关api全自动实现resnet的量化

# ====================== 核心流程 ======================

def ptq_quantize(model, calibration_data, backend="fbgemm"):
    
    #PTQ 三步走: prepare_fx → calibrate → convert_fx。
	
    # 取一个样本用于 FX trace
    example_input = calibration_data[:1]
    model.eval()

    # ---- Step 1: 插入 Observer (prepare_fx) ----
    qconfig = get_default_qconfig(backend)
    qconfig_mapping = QConfigMapping().set_global(qconfig)

    prepared_model = prepare_fx(
        model,
        qconfig_mapping,
        example_inputs=(example_input,),
    )

    # ---- Step 2: 校准 ----
    # 用一批数据跑前向传播，Observer 在这过程中收集统计信息
    # 不需要标签，不需要很多数据，256~1024 个样本即可
    prepared_model.eval()
    with torch.no_grad():
        for i in range(0, len(calibration_data), 32):
            batch = calibration_data[i:i + 32]
            prepared_model(batch)

            
    # ---- Step 3: 转换 ----
    # 根据 Observer 收集的统计信息计算 scale/zero_point
    # 将 FP32 权重量化为 INT8，插入 quantize/dequantize 节点
    print("[Step 3/3] 转换为 INT8 模型 (convert_fx)")
    quantized_model = convert_fx(prepared_model)

    print("完成!")
    return quantized_model

step 1

1 2	qconfig = get_default_qconfig(backend) qconfig_mapping = QConfigMapping().set_global(qconfig)

qconfig:根据指定的后端，获取一套默认的量化配置，规定量化的规则（INT8就在这里指定），例如权重、激活值怎么量化、使用什么observer等等
qconfig_mapping：这个 QConfig 应该应用到模型中的哪些算子（确定量化位置），set_global则是默认整个模型都用qconfig

总结：QConfig：规定“怎么量化“；QConfigMapping：规定“哪些模块使用哪个 QConfig”

prepared_model = prepare_fx(
    model,
    qconfig_mapping,
    example_inputs=(example_input,),
) #插入Observer，进行BN folding

这一步可以理解为使用torch的fx工具观测模型的具体量化结构，也就是一种可分析、可修改的计算图（例如Conv → BN → ReLU这样的结构）

在获得量化计算图、量化位置和样本输入（确认输入形状，随便输入一个随机值都可以）后，prepare_fx会在要做量化的位置插入Observer，Observer会在后续校准过程中统计min/max值用于量化计算

另外，BN folding也会在这里自动进行

自定义Q config

我们之前使用的是get_default_qconfig(backend)，即默认量化配置，现在我们进行一个自定义配置：

from torch.ao.quantization import QConfig
from torch.ao.quantization.observer import (
    MinMaxObserver,  #常规tensor ob
    PerChannelMinMaxObserver, #逐通道ob
    HistogramObserver,  #  更强，会统计分布情况，将量化范围放在主要分布区间
)
defualt_config = get_default_qconfig(backend)
My_Qconfig_1 = QConfig(
    	activation=HistogramObserver.with_args(dtype=torch.qint8),
    	weight=PerChannelMinMaxObserver.with_args(
        	qscheme=torch.per_channel_symmetric, #symmetric为对称量化
        	dtype=torch.qint8,
        	reduce_range=False,
    	)
)
qconfig_mapping = (
    QConfigMapping()
    .set_global(defualt_config)	#未说明则用默认config
    .set_object_type(nn.Conv2d, My_Qconfig_1) #卷积层用自定义
    .set_object_type(nn.BatchNorm2d, None) #None为不量化
    .set_object_type(nn.LSTM, None)
    .set_module_name("conv1", None) #也可以根据module的名字更具体的选择
    #.set_module_name_regex("layer1.*", My_Qconfig_1) #对某一层
)

"""

prepared_model = prepare_fx(
    model,
    qconfig_mapping,
    example_inputs=(example_input,),
)
.....

"""

step 2 calibration

后续要操作的模型都是prepared_model，这时需要一批能代表整个数据集的数据，来校准量化范围。这一步比较简单

step 3

已插入的Observer会自动保存统计值，convert_fx会自动根据这些值将模型转换为量化模型

总结一下

这里的量化步骤大概就是：

1	制定量化规则（这里是直接用的默认方案）-> 插入 Observer -> calibration校准 ->执行量化

DCCRN中的 BN folding难题

在DCCRN中，complex conv大概流程是：

real2real = self.real_conv(real,)
imag2imag = self.imag_conv(imag,)

real2imag = self.imag_conv(real)
imag2real = self.real_conv(imag)

real = real2real - imag2imag
imag = real2imag + imag2real
out = torch.cat([real, imag], self.complex_axis)

也就是说，对于一个卷积层，例如real_conv，其中每一个卷积核W，都同时作用于2个输入，并且这2个输出最后也分别到了不同的通道channel（r2r进入了real，i2r进入了imag，real和imag是cat关系），因此最后这2个输入在BN中也对应了2套参数，也就是说，对于有C个卷积核的conv，对应了2C个scale，而对于BN folding：

$y=aW*x+(ab_{conv}+b)=W'*x+b'$

一个W’只能融合一个a

也就是说，我们这里根本没办法直接融合！想融合的话，我们就必须再添加一倍的新conv，这个新conv仅在推理中使用，加载训练好的参数，W分开融合，但这样卷积的权重存储就会直接翻倍，这与我们量化的初衷相矛盾

因此这里的最佳思路就是不进行BN fold，那BN是否需要INT8量化呢？

对于不同的算子（Conv，linear..)，量化后端有不同的INT8 kernel，可以理解为专门为这个算子设计的“量化快速通道”，而folding过后，就可以直接使用conv 算子（因为BN参数融合进W了），这也是现在的常见推理后端倾向于的做法；而如果不folding，有些后端可能没有针对BN做专门的INT8 kernel优化，导致性能受阻；因此，

如果不folding，BN还是保持fp32比较好

在以上基础上进行QAT

PTQ是对训练好的模型直接量化，因此如果任务比较复杂时，在量化-反量化的过程中，就会产生较大的精度损失，也就是量化误差；因此，为了解决”模型从没见过量化误差”这个问题，QAT提出在训练的过程中就让模型看到量化误差，从而进行学习

可以从以下公式简单理解：

$\hat W=Dequant(Quant(W))$

在训练过程中权重仍按照FP32维护，但其在forward的过程中，就已经引入了量化带来的误差

QAT代码体现就是PTQ的流程加入到训练中：

My_Qconfig_1 = QConfig(
    	activation=HistogramObserver.with_args(dtype=torch.qint8),
    	weight=PerChannelMinMaxObserver.with_args(
        	qscheme=torch.per_channel_symmetric, #symmetric为对称量化
        	dtype=torch.qint8,
        	reduce_range=False,
    	)
)
qconfig_mapping = (
    QConfigMapping()
    .set_global(My_Qconfig_1)
)

qat_model = prepare_qat_fx( #QAT有专门的API接口
    model,
    qconfig_mapping,
    example_inputs=(example_input,)
)

optimizer = torch.optim.Adam(
    qat_model.parameters(), #使用fake q 的模型参数
    lr=1e-4
)

qat_model.train()
for x, target in train_loader:
    optimizer.zero_grad()
    output = qat_model(x)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

qat_model.eval()
quantized_model = convert_fx(qat_model)#训练完成后转换为真正的量化模型

在QAT过程中，observer会在每个batch中都不断积累更新，配合fake quant

用torch 的FX Graph Mode实现DCCRN的PTQ

要想直接调用API实现DCCRN量化有几个问题（除了上面讲到的BN folding）

FBGEMM 对 ConvTranspose 不支持使用pre-channel 进行量化
DCCRN本身的complexConv是自封装的，FX本身无法直接识别
复数LSTM对于FX来说无法直接跟踪（即使只量化Conv，FX也会跟踪完整的模型结构）

对于第一个问题，我们只需要将encoder和decoder的qconfig分开配置即可：

encoder_qconfig = get_default_qconfig(backend)

# Decoder qconfig: ConvTranspose2d 只支持 per-tensor 权重
decoder_qconfig = QConfig(
    activation=MinMaxObserver.with_args(
        dtype=torch.quint8,
        qscheme=torch.per_tensor_affine,
        reduce_range=False,
    ),
    weight=MinMaxObserver.with_args(
        dtype=torch.qint8,
        qscheme=torch.per_tensor_symmetric,
    ),
)

对于第二个问题，我们刚刚提到可以直接通过module的名字来寻找量化位置，因此我们可以：

encoder_names = []
decoder_names = []

#torch的name_modules()命名规范： 
for i in range(len(model.encoder)):		
    encoder_names.append(f"encoder.{i}.0.real_conv")
    encoder_names.append(f"encoder.{i}.0.imag_conv")

for i in range(len(model.decoder)):
    decoder_names.append(f"decoder.{i}.0.real_conv")
    decoder_names.append(f"decoder.{i}.0.imag_conv")

#全局默认不量化
qconfig_mapping = QConfigMapping().set_global(None)
#逐层映射
for name in encoder_names:
    qconfig_mapping.set_module_name(name, encoder_qconfig)
for name in decoder_names:
    qconfig_mapping.set_module_name(name, decoder_qconfig)
    

prepared_model = prepare_fx(
    model,
    qconfig_mapping,
    example_inputs=(example_input,),
)
#校准。。。
#转化
quantized_model = convert_fx(prepared_model)

后续再进行校准和量化即可

但纯调用prepare_fx，convert_fx在模型结构比较复杂时实在是很难使用，因此下一节，我们直接在模型内部进行更精细的量化

小董的BLOG

量化学习