4、模型导出实战：动手将训练好的ResNet或YOLO模型导出为ONNX，并验证导出文件的正确性

模型导出是ONNX边缘部署链路中承上启下的关键环节。前面章节我们完成了环境准备与PyTorch模型训练，本章将聚焦于“如何把训练好的PyTorch模型（以ResNet18和YOLOv5s为例）正确转换为ONNX格式”，并建立一套可复用的验证流程。导出不是简单的torch.onnx.export调用，它涉及输入尺寸固定、动态轴声明、算子兼容性检查、输出张量比对等多个维度。任何一个环节的疏漏，都可能导致部署端推理结果与训练端不一致，或在边缘设备上根本无法加载。本章将按照“导出前准备→执行导出→静态验证→动态比对→排障复盘”的完整路径展开，确保你不仅能导出文件，更能确认这个文件是“可信赖”的。

上图展示了模型导出与验证的完整工程链路。四个阶段环环相扣，其中“动态比对”是判断导出正确性的最终依据。下面我们将逐一深入每个阶段的具体操作与判断标准。

4.1 导出前准备：模型状态与输入规范的确定

导出前的准备工作决定了ONNX文件的“基因质量”。很多初学者直接对训练状态的模型执行导出，导致BN层统计量混乱、Dropout随机性被固化，最终推理结果异常。正确的做法是：先加载训练好的权重，然后调用model.eval()切换到推理模式，再构造一个与部署场景完全一致的输入张量。

核心原理：训练模式与推理模式的本质差异

PyTorch模型在训练模式下，BatchNorm层会使用当前batch的均值和方差进行归一化，同时更新全局统计量；Dropout层会随机丢弃神经元。而推理模式下，BatchNorm使用训练阶段累积的全局统计量，Dropout则变为恒等映射。如果导出时未切换到eval模式，ONNX文件中会固化训练时的随机行为或错误的归一化参数，导致部署端输出与预期严重偏离。

操作流程：输入规范的固定与动态轴声明

1. 加载模型权重：使用torch.load或model.load_state_dict恢复训练好的参数。
2. 切换eval模式：调用model.eval()，并关闭梯度计算with torch.no_grad()。
3. 构造样例输入：对于ResNet18，输入通常为torch.randn(1, 3, 224, 224)；对于YOLOv5s，输入为torch.randn(1, 3, 640, 640)。注意，这里的batch size和分辨率必须与部署目标一致。
4. 确定动态轴：如果边缘设备需要处理不同batch size或不同分辨率的输入，需要在dynamic_axes参数中声明。例如dynamic_axes={'input': {0: 'batch_size', 2: 'height', 3: 'width'}}。但需注意，动态轴会增加算子复杂度，部分边缘NPU不支持动态shape，因此建议优先固定尺寸。

4.2 执行导出：torch.onnx.export的参数精调与日志解读

导出命令本身并不复杂，但参数选择直接影响ONNX文件的兼容性和推理效率。我们以ResNet18和YOLOv5s为例，展示两种典型导出配置。

ResNet18导出示例

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
print("ResNet18 ONNX导出完成")

YOLOv5s导出示例（含预处理层）

import torch
from models.experimental import attempt_load

model = attempt_load('yolov5s.pt', map_location='cpu')
model.eval()

dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    dummy_input,
    "yolov5s.onnx",
    export_params=True,
    opset_version=12,
    do_constant_folding=True,
    input_names=['images'],
    output_names=['output'],
    dynamic_axes={'images': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
print("YOLOv5s ONNX导出完成")

关键参数解读

4.3 静态验证：使用onnxruntime与checker进行文件体检

导出完成后，不能直接拿去部署。我们需要对ONNX文件进行“体检”，包括结构完整性检查、shape推断和算子支持度分析。

判断方法：onnx.checker与shape_inference

import onnx
import onnxruntime as ort

# 1. 结构完整性检查
model = onnx.load("resnet18.onnx")
onnx.checker.check_model(model)
print("checker通过：模型结构完整")

# 2. 推断shape信息
inferred_model = onnx.shape_inference.infer_shapes(model)
onnx.save(inferred_model, "resnet18_inferred.onnx")
print("shape推断完成")

# 3. 使用onnxruntime加载并查看输入输出
sess = ort.InferenceSession("resnet18.onnx", providers=['CPUExecutionProvider'])
for inp in sess.get_inputs():
    print(f"输入: {inp.name}, shape: {inp.shape}, type: {inp.type}")
for out in sess.get_outputs():
    print(f"输出: {out.name}, shape: {out.shape}, type: {out.type}")

案例拆解：YOLOv5s导出后的静态验证结果

某次YOLOv5s导出后，执行checker时出现Error: No Op registered for Resize with domain_version of 12。分析发现，YOLOv5s中的上采样层使用了torch.nn.functional.interpolate，在opset 12中对应Resize算子，但某些边缘NPU的runtime不支持该算子的特定模式。解决方案有两种：一是将opset_version降至11，此时Resize算子行为更接近旧版；二是在导出前修改模型代码，将interpolate替换为nn.Upsample，并设置mode='nearest'。最终选择opset 11并重新导出，checker顺利通过。

4.4 动态比对：PyTorch与ONNX Runtime的输出一致性验证

静态验证只能证明文件结构正确，无法证明数值正确。动态比对是验证导出正确性的“金标准”。我们需要用同一输入分别跑PyTorch模型和ONNX Runtime，然后比较输出张量的差异。

操作流程：误差指标计算与阈值判定

import numpy as np
import torch
import onnxruntime as ort

# 准备测试输入
test_input = torch.randn(1, 3, 224, 224)

# PyTorch推理
with torch.no_grad():
    torch_output = model(test_input).numpy()

# ONNX Runtime推理
sess = ort.InferenceSession("resnet18.onnx")
ort_output = sess.run(None, {"input": test_input.numpy()})[0]

# 计算误差
cos_sim = np.dot(torch_output.flatten(), ort_output.flatten()) / (
    np.linalg.norm(torch_output) * np.linalg.norm(ort_output))
max_abs_err = np.max(np.abs(torch_output - ort_output))
mean_abs_err = np.mean(np.abs(torch_output - ort_output))

print(f"余弦相似度: {cos_sim:.6f}")
print(f"最大误差: {max_abs_err:.6f}")
print(f"平均误差: {mean_abs_err:.6f}")

# 判定标准
if cos_sim > 0.9999 and max_abs_err < 1e-4:
    print("验证通过：导出模型与原始模型输出高度一致")
else:
    print("验证失败：请检查导出参数或模型结构")

案例拆解：YOLOv5s输出差异的定位与修复

某次YOLOv5s导出后，动态比对发现最大误差达到0.35，远超阈值。通过逐层对比，定位到问题出在Focus模块。YOLOv5s的Focus层将输入按空间位置切片后拼接，在PyTorch中实现为tensor.view和cat操作。但ONNX导出时，该操作被转换为Slice和Concat的组合，由于opset版本差异，Slice的步长处理方式不同，导致输出错位。解决方案：在导出前将Focus层替换为等价的Conv2d（stride=2），避免切片操作。修改后重新导出，最大误差降至1e-6，验证通过。

4.5 常见误区与工程复盘

在模型导出实战中，以下误区最容易导致部署失败或推理异常，需要特别警惕。

• 误区一：忽略预处理对齐。训练时的归一化参数（mean/std）必须在部署端复现。ONNX模型只包含网络结构，不包含预处理逻辑。如果部署端使用不同mean/std，输出将完全错误。
• 误区二：动态轴滥用。虽然ONNX支持动态shape，但边缘NPU通常需要固定shape以优化内存布局。动态轴会导致推理引擎无法预分配内存，性能下降30%~50%。
• 误区三：忽略opset版本兼容性。边缘设备的runtime可能只支持opset 10或11。导出前应查阅目标设备的算子支持列表，选择最合适的opset版本。
• 误区四：只验证单一样本。建议使用至少10张不同风格的测试图片，覆盖不同光照、分辨率和目标数量，确保导出模型在各种输入下均表现一致。

工程复盘：一次完整的导出排障记录

某边缘盒子项目需要部署YOLOv5s，目标设备为瑞芯微RK3588 NPU。导出时使用opset 12，checker通过，但NPU加载时报错“Unsupported op: Resize”。分析过程：首先确认RK3588的rknn-toolkit2支持opset 10~12，但Resize算子的坐标变换模式（coordinate_transformation_mode）仅支持“asymmetric”和“half_pixel”。检查ONNX模型，发现Resize节点使用了“align_corners=True”，对应模式为“align_corners”。解决方案：在导出前修改模型代码，将interpolate的align_corners参数改为False，并设置mode='nearest'。重新导出后，NPU加载成功，推理速度达到35ms/帧，满足实时性要求。复盘经验：导出前必须查阅目标NPU的算子支持文档，不能仅依赖onnxruntime的兼容性。