第4章：模型导出与量化——从FP32到INT8的实战之路

模型导出和量化，说白了就是把训练好的“大胖子”模型，压缩成能在芯片上跑起来的“精干小伙”。我刚开始接触这块时，也觉得不就是转个格式嘛，结果踩了不少坑。今天咱们就把PyTorch/TensorFlow模型导出为ONNX、INT8量化原理、校准数据集制作、量化精度损失分析这几个环节，一个一个掰开揉碎了讲。

4.1 模型导出：PyTorch/TensorFlow → ONNX

ONNX（Open Neural Network Exchange）是模型交换的中间格式。你想想看，训练用PyTorch，部署用TensorRT，中间总得有个“翻译官”。ONNX就是这个角色。

4.1.1 PyTorch导出ONNX

PyTorch导出ONNX主要用 torch.onnx.export。我个人习惯在导出前先做一次推理验证，确保模型是正常的。

import torch
import torch.onnx

# 假设model是你的ISP模型，dummy_input是示例输入
model.eval()
dummy_input = torch.randn(1, 3, 640, 640)  # batch=1, 3通道, 640x640

torch.onnx.export(
    model,
    dummy_input,
    "isp_model.onnx",
    export_params=True,        # 导出训练好的参数
    opset_version=11,          # 推荐11或更高
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX导出完成")

4.1.2 TensorFlow导出ONNX

TensorFlow导出ONNX需要借助 tf2onnx 工具。流程稍微绕一点，但原理一样。

import tf2onnx
import tensorflow as tf

# 加载你的TensorFlow模型
model = tf.saved_model.load("path/to/saved_model")

# 指定输入签名
spec = (tf.TensorSpec((None, 640, 640, 3), tf.float32, name="input"),)

# 导出ONNX
output_path = "isp_model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model, input_signature=spec, opset=13, output_path=output_path
)
print(f"ONNX模型已保存至: {output_path}")

4.2 INT8量化原理：为什么是INT8？

FP32模型在AI芯片上跑，功耗高、带宽大。INT8量化就是把FP32的权重和激活值，映射到INT8的范围内（-128到127）。

核心公式很简单：

量化: Q = round(R / S) + Z
反量化: R = (Q - Z) * S

其中：
R = 真实浮点值
Q = 量化后的整数值
S = 缩放因子 (scale)
Z = 零点偏移 (zero_point)

你想想看，FP32有4个字节，INT8只有1个字节。模型体积直接缩小4倍，推理速度提升2-4倍。在ISP场景下，这个加速效果非常明显。

4.2.1 对称量化 vs 非对称量化

4.3 校准数据集制作：量化前的关键一步

量化需要“校准数据集”来确定S和Z。这个数据集不需要很大，但必须能代表真实数据分布。

我一般这样制作：

1. 从训练集中随机抽取：500-1000张图片，覆盖不同光照、场景
2. 预处理对齐：确保预处理方式和训练时完全一致
3. 保存为特定格式：推荐使用numpy数组或TFRecord，方便快速加载
4. 验证多样性：检查是否包含极端情况（过曝、欠曝、噪声大等）

import numpy as np
import cv2
import os

def make_calibration_dataset(image_dir, output_path, size=(640, 640)):
    """
    制作校准数据集
    """
    images = []
    for fname in os.listdir(image_dir)[:500]:  # 取500张
        img = cv2.imread(os.path.join(image_dir, fname))
        img = cv2.resize(img, size)
        img = img.astype(np.float32) / 255.0  # 归一化到[0,1]
        images.append(img)
    
    calib_data = np.stack(images, axis=0)  # [N, H, W, C]
    np.save(output_path, calib_data)
    print(f"校准数据集已保存: {calib_data.shape}")
    return calib_data

4.4 量化精度损失分析：到底损失了多少？

量化一定会损失精度，关键是要控制在可接受范围内。我通常用以下指标评估：

4.4.1 常用评估指标

4.4.2 精度损失常见原因

• 权重分布不均匀：某些层数值范围过大，量化后信息丢失严重
• 激活值有异常点：个别极大值拉高了整个量化范围
• 校准数据集不匹配：校准数据分布和实际推理数据分布不一致
• 敏感层未特殊处理：某些层（如第一层、最后一层）对量化特别敏感

4.5 知识结构总览

下面这张图总结了本章的核心流程，从模型导出到量化部署，每一步都环环相扣。

4.6 实战建议与避坑总结

最后，我把这些年积累的经验浓缩成几条建议：

• 导出ONNX后一定要验证：用onnxruntime跑一遍，对比输出是否一致。我见过太多导出后算子不对齐的情况。
• 量化前先做敏感层分析：逐层量化，找出哪些层对精度影响最大。这些层可以考虑保留FP32。
• 校准数据集宁多勿少：500张是底线，1000张更稳妥。场景覆盖越全，量化后泛化能力越好。
• 量化后做全量测试：不要只看几个指标，要在完整测试集上跑一遍。我遇到过量化后PSNR没掉，但某个特定场景全黑的情况。