4、模型转换与量化：将训练好的模型（PyTorch/TF）转换为TFLite/ONNX格式，以及INT8量化实战

模型训练完了，在电脑上跑得挺欢。但你要把它塞进一个只有几百KB内存的MCU里？嗯，这里就有故事了。

我刚开始做边缘端部署时，犯过一个低级错误——直接把PyTorch的.pth文件往STM32里拷。结果可想而知，芯片根本不认识这玩意儿。后来才明白，模型从训练框架到推理引擎，中间必须经过一道「翻译」工序。这就是模型转换。

而量化，说白了就是给模型「瘦身」。把32位浮点数变成8位整数，模型体积直接缩到四分之一，推理速度翻倍。代价呢？精度会掉一点点。但在我做过的项目中，大部分场景下这点精度损失完全能接受。

4.1 为什么需要模型转换与量化

你想想看，PyTorch和TensorFlow是训练框架，它们的设计目标是灵活、可微分、能反向传播。而TFLite和ONNX是推理格式，它们的目标是轻量、高效、能跑在特定硬件上。

这两者的基因完全不同。所以我们需要一座桥。

我记得有个项目，模型在GPU上推理只要5ms，但转到Cortex-M4上后，一次推理花了3秒。后来做了INT8量化加模型剪枝，才压到200ms以内。这就是量化的价值。

4.2 PyTorch模型转ONNX

ONNX是个好东西。它像一个通用翻译器，能把PyTorch、TF、PaddlePaddle的模型都转成同一种中间格式。然后你再从ONNX转到目标平台的格式。

我个人习惯先用ONNX做一次中转，因为它的生态最成熟，踩坑的人也最多，网上资料好找。

下面是一个典型的PyTorch转ONNX示例：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = torch.load('model.pth')
model.eval()

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

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    'model.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'}
    }
)

print("ONNX导出成功！")

4.3 TensorFlow模型转TFLite

TF转TFLite相对简单，因为都是Google自家的东西。但要注意，TF1和TF2的转换API完全不同。现在主流都用TF2了。

我曾经在TF1时代被坑过——SavedModel和GraphDef两种格式搞混，折腾了一整天。后来学乖了，统一用SavedModel。

import tensorflow as tf

# 加载SavedModel模型
model = tf.saved_model.load('saved_model_dir')

# 创建转换器
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model_dir')

# 如果不做量化，直接转换
tflite_model = converter.convert()

# 保存TFLite模型
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

print("TFLite转换成功！")

这里有个坑：如果你的模型里有自定义算子，TFLite默认不支持。需要注册自定义算子，或者用TF Select算子。嗯，这个后面会细讲。

4.4 INT8量化实战

量化才是重头戏。INT8量化分两种：训练后量化和量化感知训练。

训练后量化最简单，模型训练完直接转。适合大部分场景。量化感知训练则是在训练过程中模拟量化效果，精度更高，但需要改训练代码。

我个人建议：先试训练后量化。如果精度掉得太多，再上量化感知训练。别一上来就搞复杂的。

4.4.1 TFLite INT8量化

import tensorflow as tf

# 加载模型
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model_dir')

# 设置量化方式为INT8
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 提供代表性数据集（用于校准）
def representative_dataset():
    for _ in range(100):
        data = tf.random.normal([1, 224, 224, 3])
        yield [data]

converter.representative_dataset = representative_dataset

# 指定目标类型为INT8
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# 设置输入输出类型
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

# 转换
tflite_quant_model = converter.convert()

# 保存
with open('model_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print("INT8量化模型生成成功！")

4.4.2 ONNX转INT8量化

ONNX的量化工具链稍微复杂一些。我常用的是ONNX Runtime的量化工具：

from onnxruntime.quantization import quantize_dynamic, QuantType

# 动态量化（不需要校准数据）
quantize_dynamic(
    'model.onnx',
    'model_quant.onnx',
    weight_type=QuantType.QInt8
)

print("ONNX动态量化完成！")

如果需要静态量化（精度更高），需要校准数据：

from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization import CalibrationMethod

# 静态量化需要校准数据
calibration_data = ...  # 你的校准数据集

quantize_static(
    'model.onnx',
    'model_quant.onnx',
    calibration_data_reader=calibration_data,
    quant_format=QuantType.QInt8,
    per_channel=True,
    calibration_method=CalibrationMethod.MinMax
)

4.5 量化后的精度验证

量化完别急着部署。先跑一遍验证集，看看精度掉了多少。

我一般定个标准：Top-1准确率下降不超过1%就算合格。如果超过2%，就要考虑用量化感知训练了。

量化方法	模型大小	推理速度	精度损失	适用场景
FP32（未量化）	100%	1x	0%	服务器端
FP16	50%	1.5-2x	~0.1%	GPU/高端NPU
INT8（训练后量化）	25%	3-4x	0.5-2%	大部分边缘设备
INT8（量化感知训练）	25%	3-4x	0.1-0.5%	精度敏感场景

4.6 常见坑与避坑指南

做模型转换和量化这几年，我踩过的坑能写一本书。挑几个典型的说说：

• 算子不支持：有些PyTorch的高级算子（如torch.einsum）ONNX不支持。解决办法是用基础算子重写，或者注册自定义算子。
• 动态shape问题：边缘设备通常要求固定输入尺寸。如果你的模型支持动态batch，记得在导出时固定下来。
• 量化后精度骤降：我曾经遇到一个模型，量化后准确率从95%掉到60%。查了半天，发现是某个层的激活值分布太广，量化时信息丢失严重。后来加了clip操作才解决。
• 端侧推理速度没提升：INT8量化理论上快4倍，但如果你用的MCU没有硬件INT8加速，实际速度可能只快1.5倍。嗯，这个要提前查芯片手册。

最后说一句：模型转换和量化不是一锤子买卖。你可能需要反复尝试不同的量化策略、校准数据集大小、算子选择。别怕麻烦，这是嵌入式AI工程师的日常。