3. 模型转换与量化基础：从PyTorch/TensorFlow到ONNX、INT8量化原理、校准数据集准备

好，咱们进入第三章。这一章我打算聊聊模型转换和量化。说实话，这两个东西是嵌入式部署的「入场券」。你模型训练得再漂亮，落不了地，那就是空中楼阁。我自己刚入行那会儿，就吃过这个亏——模型在GPU上跑得飞起，一上板子直接崩了。后来才发现，是转换环节出了岔子。

3.1 为什么需要模型转换？

你想想看，PyTorch和TensorFlow是给训练设计的。它们灵活、动态，方便你调参。但嵌入式设备要的是「死」的、固定的计算图。说白了，就是要把训练框架的「花架子」拆掉，变成一套纯粹的数学运算指令。

ONNX（Open Neural Network Exchange）就是干这个的。它像一个中间语言。你用PyTorch写了个模型，导出成ONNX。再用TensorFlow写一个，也能导出成ONNX。然后，你的嵌入式推理引擎（比如ONNX Runtime、TensorRT、NCNN）就可以直接吃ONNX了。

我个人习惯，所有模型落地前，先过一遍ONNX。为什么？因为ONNX的计算图是标准化的，很多工具都能解析它。一旦卡住，排查问题也方便。

3.2 从PyTorch到ONNX

PyTorch导出ONNX，核心就一个函数：torch.onnx.export()。但这里面的坑，不少。

我给你看个典型例子：

import torch
import torchvision.models as models

# 加载一个预训练的ResNet18
model = models.resnet18(pretrained=True)
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里我重点说几个参数：

• opset_version：这个很关键。版本太低，有些新算子不支持。版本太高，你的推理引擎可能又不认识。我一般用11或13，兼容性比较好。
• dynamic_axes：允许动态batch。你想想看，如果每次推理的batch size都固定为1，那还好。但如果你要支持不同大小的输入，就必须打开这个选项。我在项目中遇到过，没设dynamic_axes，结果换了一张分辨率不同的图片，直接报shape不匹配。
• do_constant_folding：建议打开。它会把一些固定的计算提前算好，比如BN层的参数融合。能省一点是一点。

3.3 从TensorFlow到ONNX

TensorFlow这边，稍微麻烦一点。因为TensorFlow有1.x和2.x两个版本，导出方式不同。

对于TensorFlow 2.x，推荐用tf2onnx这个工具。安装很简单：

pip install tf2onnx

然后命令行转换：

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 13

或者用Python API：

import tf2onnx
import tensorflow as tf

# 加载SavedModel
model = tf.saved_model.load('./saved_model')

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

# 转换
output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_function(
    model.signatures['serving_default'],
    input_signature=spec,
    opset=13,
    output_path=output_path
)

嗯，这里要注意。TensorFlow的图里经常有一些自定义算子，比如一些图像预处理操作。这些算子ONNX不一定认识。我建议，把预处理逻辑单独拿出来，在CPU上做，只把纯推理部分丢给ONNX。这样能省去很多麻烦。

3.4 INT8量化原理

模型转换完了，接下来就是量化。为什么需要量化？说白了，就是为了快和省。

你想想看，FP32的模型，一个参数占4个字节。INT8的模型，一个参数只占1个字节。模型体积直接缩小4倍。而且，INT8的矩阵乘法，在硬件上比FP32快得多。很多嵌入式芯片，比如高通、联发科的DSP，都对INT8做了专门优化。

量化的核心思想，就是用更少的bit来表示一个数值范围。FP32能表示的范围是 [-3.4e38, 3.4e38]，精度很高。INT8只能表示 [-128, 127] 这256个整数。所以，我们需要做一个映射。

这个映射，通常用「缩放因子」和「零点」来描述：

real_value = scale * (quantized_value - zero_point)

其中，scale 是缩放因子，zero_point 是零点偏移。这两个参数，就是量化要确定的。

常见的量化方式有两种：

• 对称量化：零点固定为0。适合权重这种正负对称的分布。
• 非对称量化：零点可以偏移。适合激活值这种非负的分布。

我个人更倾向于非对称量化，因为它能更好地利用INT8的表示范围。比如ReLU的输出全是非负数，用对称量化会浪费一半的表示能力。

3.5 校准数据集准备

量化不是随便找个范围就行的。你需要知道，你的模型在真实数据上，每一层的激活值分布是什么样的。这个「摸底」的过程，就叫校准（Calibration）。

校准需要准备一个数据集。这个数据集不需要很大，但要有代表性。我一般从训练集中抽100-500张图片，覆盖各种场景。

校准的过程，说白了就是跑一遍推理，记录每一层激活值的最大值和最小值。然后根据这些统计值，计算出最佳的scale和zero_point。

常见的校准方法有：

方法	原理	特点
MinMax	直接取激活值的全局最小/最大值	简单，但对异常值敏感
KL散度	寻找一个阈值，使得量化前后的分布差异最小	效果较好，NVIDIA TensorRT默认使用
百分比	去掉一定比例的极端值（如99.9%分位）	折中方案，鲁棒性较好

我建议，如果时间允许，三种方法都试一下，看哪个在验证集上精度最高。我在项目中遇到过，MinMax方法因为一张过曝的图片，导致量化后精度暴跌。换成KL散度就好了。

3.6 量化实战：PyTorch静态量化

来，我们看一个完整的PyTorch静态量化例子。静态量化，就是先校准，再量化。推理时不再调整参数。

import torch
import torch.quantization as quant

# 1. 准备模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 融合BN和Conv层（减少量化误差）
model = quant.fuse_modules(model, [['conv1', 'bn1', 'relu']])

# 3. 设置量化配置
model.qconfig = quant.get_default_qconfig('fbgemm')

# 4. 准备校准
model = quant.prepare(model, inplace=True)

# 5. 校准：用校准数据集跑几遍
def calibrate(model, data_loader, num_batches=10):
    model.eval()
    with torch.no_grad():
        for i, (inputs, _) in enumerate(data_loader):
            if i >= num_batches:
                break
            model(inputs)

calibrate(model, calibration_loader)

# 6. 转换到量化模型
quantized_model = quant.convert(model, inplace=False)

# 7. 保存
torch.save(quantized_model.state_dict(), 'resnet18_quantized.pth')

这里有几个关键点：

• fuse_modules：把Conv+BN+ReLU融合成一个算子。这样量化时误差更小。我建议能融合的都融合掉。
• qconfig：量化配置。fbgemm是Facebook针对x86优化的后端。如果你在ARM上跑，可以用qnnpack。
• 校准的batch数：我一般用10-20个batch就够了。太多也没意义，反而浪费时间。

好了，这一章就到这里。下一章，我们会深入推理引擎的底层，看看ONNX Runtime和TensorRT到底是怎么把模型跑起来的。到时候，我会分享一些我自己调优时踩过的坑，保证让你少走弯路。