📚 课程目录第4章模型导出

4、模型导出：从PyTorch/TensorFlow导出ONNX及注意事项

在TinyML端侧部署的完整链路中，模型导出是连接“训练”与“推理优化”的桥梁。无论你使用PyTorch还是TensorFlow训练模型，最终都需要将其转换为一种中间表示，以便后续进行量化、剪枝和转换为TFLite/ONNX Runtime等端侧推理格式。ONNX（Open Neural Network Exchange）正是这一生态中的“通用语言”。本章将深入讲解从PyTorch和TensorFlow/Keras导出ONNX的标准流程、核心参数，以及导出过程中最容易踩坑的动态轴、算子兼容性和输入输出命名问题。

上图展示了本章的核心内容框架：PyTorch与TensorFlow两条导出路径汇聚到ONNX，再围绕动态轴、算子支持和命名规范三个关键决策点展开。下面我们逐一深入。

4.1 从PyTorch导出ONNX：torch.onnx.export深度解析

核心原理

torch.onnx.export 的本质是符号跟踪（symbolic tracing）。它通过运行一次模型的前向传播（forward pass），记录下所有张量操作的执行顺序，并将这些操作映射为ONNX算子。这意味着你不需要手动构建计算图，但必须提供一个形状正确的示例输入（dummy input），因为跟踪过程依赖于实际的数据流。

导出的关键参数包括：

model：要导出的PyTorch模型，需处于eval()模式。
args：示例输入张量，其形状决定了ONNX模型的静态输入形状。
input_names / output_names：为输入输出张量指定名称，后续推理时按名称索引。
dynamic_axes：指定哪些维度是动态的，例如batch维度。
opset_version：ONNX算子集版本，影响可用算子的种类和精度。

操作流程

模型准备：调用 model.eval() 并确保所有参数已加载完毕。
构造示例输入：使用 torch.randn 生成符合输入尺寸的假数据。
执行导出：调用 torch.onnx.export，指定文件路径、示例输入、名称和动态轴。
验证导出结果：使用 onnx.checker.check_model 检查模型合法性，并用 onnxruntime 进行推理对比。

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv(x)))
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = SimpleCNN().eval()
dummy_input = torch.randn(1, 3, 32, 32)

torch.onnx.export(
    model,
    dummy_input,
    "simple_cnn.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=13
)
print("导出完成")

经验提醒：导出前务必调用 model.eval()。如果忘记切换，BatchNorm和Dropout层的行为会不一致，导致导出的ONNX模型在推理时产生错误结果。

案例拆解：动态batch的工业质检模型

某工厂部署的PCB缺陷检测模型，训练时使用固定batch=32，但端侧设备由于内存限制，推理时batch可能为1或4。若直接导出静态模型，每次改变batch都需要重新导出。解决方案是使用dynamic_axes将batch维度设为动态：

dynamic_axes = {
    "input": {0: "batch_size"},
    "output": {0: "batch_size"}
}

导出后，ONNX模型接受任意batch大小的输入。但需要注意，动态轴会增加推理引擎的内存管理开销，且某些端侧加速器（如NPU）可能不支持动态维度。因此，在TinyML场景下，建议优先使用固定batch=1的静态模型，除非业务确实需要动态batch。

4.2 从TensorFlow/Keras导出ONNX：tf2onnx实战

核心原理

tf2onnx 是一个独立的转换工具，它将TensorFlow的GraphDef或Keras模型转换为ONNX格式。与PyTorch的跟踪机制不同，tf2onnx直接解析TensorFlow的计算图结构，逐算子映射到ONNX。对于Keras模型，推荐使用 tf2onnx.convert.from_keras 方法，它内部会先冻结模型权重，再执行转换。

操作流程

安装工具：pip install tf2onnx onnxruntime
加载Keras模型：使用 tf.keras.models.load_model 加载训练好的模型。
执行转换：调用 from_keras 并指定输出路径和opset版本。
验证：使用onnxruntime加载转换后的模型，对比Keras模型的输出。

import tf2onnx
import tensorflow as tf

# 加载Keras模型
model = tf.keras.models.load_model("mobilenet_v2.h5")

# 转换
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    opset=13,
    output_path="mobilenet_v2.onnx"
)
print("转换完成")

常见问题与处理

问题现象	可能原因	解决方案
转换报错：Unsupported Op	Keras层使用了tf2onnx不支持的算子	升级opset版本；或使用 `tf2onnx.rewriter` 自定义映射
转换后精度下降	模型包含训练特有的算子（如Dropout）	确保模型处于推理模式（`model.trainable = False`）
输入形状不匹配	Keras模型输入未指定batch维度	使用 `model.input_shape` 检查，必要时重新定义输入层

风险提示：tf2onnx对TensorFlow 2.x的支持已趋于稳定，但若模型包含自定义层（如自定义损失或自定义激活），转换时可能失败。此时需要先使用 tf.function 将自定义层包装为标准算子，或考虑改用PyTorch重新实现模型。

4.3 导出时的三大注意事项

4.3.1 动态轴：灵活性与性能的权衡

动态轴允许ONNX模型接受可变大小的输入，但代价是推理引擎无法预先分配固定内存，可能降低执行效率。在TinyML场景中，强烈建议将batch维度固定为1，因为端侧设备通常处理单帧数据。如果必须支持序列长度变化（如语音特征），可以只将序列维度设为动态，但需确认目标推理引擎（如ONNX Runtime Mobile）是否支持。

4.3.2 算子支持：opset版本与兼容性

ONNX算子集（opset）不断演进，新版本会添加新算子或修改旧算子行为。选择opset版本时需考虑：

端侧推理引擎的支持范围：例如，ONNX Runtime Mobile通常支持opset 13-18，但某些NPU加速器可能只支持opset 11。
算子映射的完整性：较新的opset可能包含更精确的算子定义，但也会增加转换工具的负担。
自定义算子的处理：如果模型包含PyTorch/TF中没有ONNX对应算子的操作，需要注册自定义算子或改用等价的标准算子组合。

# 检查ONNX模型支持的opset版本
import onnx
model = onnx.load("model.onnx")
print("IR版本:", model.ir_version)
print("Opset版本:", model.opset_import[0].version)

4.3.3 输入输出命名：端侧推理的索引依据

ONNX模型中的输入输出名称是推理时获取数据的唯一标识。在端侧部署时，你需要根据这些名称来填充输入张量和读取输出结果。命名不规范（如使用默认的"input.1"、"output.1"）会导致后续代码可读性差，且容易出错。

最佳实践：

使用语义化名称，如 "image"、"scores"、"boxes"。
保持名称唯一，避免与ONNX内部节点名冲突。
在导出时显式指定 input_names 和 output_names，不要依赖默认值。

  实战建议：在导出后，使用 onnxruntime 进行一次推理验证，打印输入输出的名称和形状，确保与预期一致。这能避免在端侧调试时才发现名称不匹配的问题。

4.4 常见误区与总结复盘

误区	后果	正确做法
忘记调用 `model.eval()`	BatchNorm统计量错误，推理精度下降	导出前显式设置 `eval()` 模式
使用训练时的batch大小导出	端侧内存不足或推理延迟高	使用batch=1导出，或设置动态轴
忽略opset版本兼容性	端侧推理引擎无法加载模型	查阅目标引擎文档，选择合适opset
输入输出命名随意	端侧代码难以维护，易出错	使用语义化名称并记录在部署文档中

总结：模型导出是TinyML部署中看似简单却影响深远的环节。PyTorch用户应熟练掌握 torch.onnx.export 的跟踪机制和动态轴配置；TensorFlow用户则需熟悉 tf2onnx 的转换流程和常见算子兼容性问题。无论哪条路径，都必须将动态轴、算子支持和命名规范作为导出前的“三查”项。一个规范的ONNX模型，能为后续的量化、剪枝和端侧转换节省大量调试时间。建议在项目初期就建立统一的导出脚本和验证流程，确保每次导出的模型都是可复现、可追踪的。