4. 模型训练与量化感知训练（QAT）

模型训练完了，浮点精度看着不错，一部署到芯片上就崩了。这种事我遇到过好几次。说白了，FP32的模型在PC上跑得飞起，但到了ISP的NPU或者DSP上，内存带宽和算力根本扛不住。所以我们必须做量化——把FP32转成INT8。

但直接转，精度掉得厉害。这时候就需要量化感知训练（QAT）来救场。这一章，我把整个流程拆开讲清楚。

4.1 训练流程：从FP32到INT8的必经之路

我个人习惯把流程分成三个阶段：

1. FP32预训练：先训一个高精度的浮点模型，这是基线。
2. 插入伪量化节点：在训练图中插入FakeQuantize操作，模拟量化误差。
3. 微调（QAT）：带着伪量化节点继续训几个epoch，让权重和激活值适应INT8的表示范围。

你想想看，如果不做QAT，直接PTQ（Post-Training Quantization），对于降噪模型这种对细节敏感的任务，PSNR可能直接掉2-3个dB。图像上就是一片一片的色块和伪影。

4.2 FP32转INT8的量化方法

量化说白了就是用一个低精度的数去近似一个高精度的数。对于INT8，我们通常用对称量化或非对称量化。

量化方法	公式	适用场景	我踩过的坑
对称量化	q = round(x / scale)	权重（分布接近0对称）	ReLU后的激活值用对称量化，浪费了负半轴
非对称量化	q = round((x - zero_point) / scale)	激活值（分布偏移）	zero_point计算不当，精度反而更差

对于降噪模型，我建议权重用对称量化，激活值用非对称量化。为什么？因为降噪模型的权重经过BN和ReLU后，分布往往偏向正数，非对称量化能更好地利用INT8的256个bin。

4.3 QAT在降噪模型中的应用

QAT的核心思想很简单：在训练前向图中插入FakeQuantize节点，让模型在训练时就“感受”到量化误差。反向传播时，使用直通估计器（STE）来近似量化函数的梯度。

我曾经在一个夜景降噪项目里，直接PTQ后PSNR掉了2.8dB。后来改成QAT，只微调了5个epoch，PSNR只掉了0.3dB。效果立竿见影。

4.3.1 PyTorch实现示例

PyTorch从1.8开始提供了torch.quantization模块，做QAT很方便。下面是一个典型的流程：

import torch
import torch.nn as nn
import torch.quantization as quant

# 1. 定义模型
class DenoiseNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 3, 3, padding=1)
        # 插入伪量化节点
        self.quant = quant.QuantStub()
        self.dequant = quant.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)  # 量化输入
        x = self.relu(self.conv1(x))
        x = self.conv2(x)
        x = self.dequant(x)  # 反量化输出
        return x

# 2. 准备QAT
model = DenoiseNet()
model.train()
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
quant.prepare_qat(model, inplace=True)

# 3. 正常训练（优化器、损失函数不变）
for epoch in range(5):
    for data, target in dataloader:
        output = model(data)
        loss = nn.MSELoss()(output, target)
        loss.backward()
        optimizer.step()

# 4. 转换到INT8
model.eval()
model = quant.convert(model, inplace=True)

4.3.2 TensorFlow实现示例

TensorFlow这边，推荐用tf.quantization.quantize_model或者TFLite的QAT API。我个人更常用TFLite的QuantizationAwareTraining：

import tensorflow as tf
import tensorflow_model_optimization as tfmot

# 1. 定义模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, padding='same', input_shape=(None, None, 3)),
    tf.keras.layers.ReLU(),
    tf.keras.layers.Conv2D(3, 3, padding='same')
])

# 2. 应用QAT
qat_model = tfmot.quantization.keras.quantize_model(model)

# 3. 编译并训练
qat_model.compile(optimizer='adam', loss='mse')
qat_model.fit(train_dataset, epochs=5)

# 4. 转换为TFLite INT8
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

4.4 QAT的关键参数与调优

QAT不是简单地插入节点就完事。有几个参数直接影响最终精度：

• 量化粒度：逐层量化 vs 逐通道量化。逐通道量化精度更高，但硬件支持有限。
• 校准数据集：用于统计激活值的范围。我一般从训练集里抽500-1000张图，覆盖不同噪声水平。
• 微调epoch数：3-10个epoch足够。太多会过拟合到量化噪声上。

我记得有一次，校准集只用了几十张图，结果量化后的模型在暗光场景下全黑了。后来换成1000张覆盖全场景的图，问题解决。

4.5 知识体系流程图

下面这张图总结了从FP32训练到INT8部署的完整链路：

4.6 实战中的避坑指南

最后，分享几个我亲身踩过的坑：

• 校准集太单一：只用明亮场景的图，结果暗光下模型输出全是噪点。后来我混合了不同ISO、不同光照的图，问题解决。
• BN层融合顺序：QAT前一定要先fuse BN和Conv，否则量化误差会放大。PyTorch里用torch.quantization.fuse_modules。
• 硬件不支持逐通道量化：有些NPU只支持逐层量化。这时候需要手动调整量化参数，或者用per_channel=False。