4、模型导出与格式转换：ONNX简介、PyTorch模型转ONNX、ONNX模型验证与可视化（Netron）

模型训练完了，然后呢？

说实话，很多同学在训练阶段跑得飞起，一到部署就卡壳。原因很简单——训练框架和推理引擎之间，缺了一座桥。这座桥，就是ONNX。

今天我们就来聊聊这座桥怎么搭。我会从ONNX是什么讲起，然后手把手带你走一遍PyTorch转ONNX的流程，最后教你用Netron看一眼模型长啥样。

4.1 ONNX到底是什么？

ONNX，全称Open Neural Network Exchange。说白了，它就是一套模型中间表示标准。

你想想看，PyTorch有自己的一套图结构，TensorFlow也有自己的。各家各玩各的，部署的时候就得写一堆适配代码。ONNX的出现，就是为了统一这个局面。

它的核心价值就两点：

• 跨框架互通：PyTorch训练的模型，转成ONNX后，可以直接在TensorRT、OpenVINO、ONNX Runtime上跑。
• 推理优化：ONNX本身支持图优化、算子融合、量化等操作。我见过一个项目，转成ONNX后用ONNX Runtime跑，推理速度直接快了30%。

我个人习惯把ONNX比作「通用语言」。你用中文写文章（PyTorch），我用法语写（TensorFlow），互相看不懂。但如果我们都翻译成世界语（ONNX），那谁都能读。

4.2 PyTorch模型转ONNX

这一步其实不复杂，但坑不少。我先给你看标准流程，再讲避坑指南。

4.2.1 标准转换流程

PyTorch官方提供了torch.onnx.export这个接口。用法如下：

import torch
import torchvision.models as models

# 1. 加载模型，设为eval模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出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'}
    }
)
print("导出成功！")

这里有几个参数我要重点说一下：

• opset_version：我建议用11或更高。版本太低，有些新算子不支持。版本太高，老推理引擎可能不兼容。11是个比较稳妥的选择。
• dynamic_axes：这个很关键。如果你希望模型支持动态batch size，就得配这个参数。否则导出的模型只能跑固定batch。
• do_constant_folding：建议打开。它会把一些固定计算提前算好，减少推理时的计算量。

4.2.2 避坑指南

我曾经在一个项目里卡了整整两天，就为了一个算子转换问题。这里把常见坑列出来：

• 动态控制流：如果你的模型里有if、for这种动态控制流，ONNX可能不支持。解决办法是用torch.where或者torch.stack替代。
• 自定义算子：自己写的torch.autograd.Function，ONNX不认识。需要注册自定义算子，或者用torch.onnx.symbolic做映射。
• 输入输出类型：ONNX不支持torch.long以外的整数类型。如果你的输入是torch.int32，记得转一下。
• BatchNorm和Dropout：导出前一定要调用model.eval()。否则BN层会使用训练时的统计量，结果完全不对。

4.3 ONNX模型验证

导出完了，怎么知道对不对？

我的做法是：用原始PyTorch模型和ONNX模型，喂同样的输入，对比输出差异。

import onnx
import onnxruntime
import numpy as np

# 1. 检查ONNX模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过！")

# 2. 用ONNX Runtime跑推理
ort_session = onnxruntime.InferenceSession("resnet18.onnx")

# 构造输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 跑推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(None, ort_inputs)

# 3. 对比PyTorch结果
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data))

# 计算差异
diff = np.abs(ort_outputs[0] - torch_output.numpy()).max()
print(f"最大差异：{diff:.6f}")

if diff < 1e-4:
    print("✅ 精度验证通过！")
else:
    print("❌ 精度有问题，请检查模型！")

这里有个经验值：最大差异在1e-4以内，基本没问题。如果超过1e-3，就要警惕了。我遇到过因为opset_version选错，导致某些算子精度下降的情况。

4.4 用Netron可视化模型

代码验证只能保证数值正确，但模型结构对不对，还得靠眼睛看。

Netron就是干这个的。它是个轻量级的模型可视化工具，支持ONNX、TensorFlow、PyTorch等多种格式。

怎么用？

• 在线版：打开浏览器，访问netron.app，直接把.onnx文件拖进去。
• 离线版：pip install netron，然后netron resnet18.onnx，会自动在浏览器打开。

打开后你会看到什么？

• 左侧是模型的计算图，每个节点代表一个算子。
• 点击节点，右侧会显示该算子的详细信息，比如输入输出形状、属性参数。
• 顶部可以搜索节点，方便定位特定层。

我个人习惯用Netron做三件事：

1. 检查输入输出形状：确保和预期一致。特别是动态batch的配置，一眼就能看出来。
2. 确认算子类型：看看有没有出现Unknown或者Identity这种奇怪的算子。如果有，说明转换有问题。
3. 查看图结构：有时候模型里多了些没用的节点，比如Shape、Gather，这些会影响推理性能。Netron上一目了然。

4.5 总结一下

模型导出这件事，说难不难，说简单也不简单。核心就三步：

• 转格式：用torch.onnx.export，注意opset版本和dynamic_axes。
• 验精度：对比PyTorch和ONNX Runtime的输出，差异控制在1e-4以内。
• 看结构：用Netron可视化，确认算子、形状、图结构都没问题。

这三步走完，你的模型就准备好进入下一阶段了——量化、剪枝、或者直接上推理引擎。嗯，这些我们后面章节再聊。