3. YOLO模型基础与导出：从PyTorch训练到ONNX
这一章，我们聊聊模型部署的起点。说白了，就是怎么把你在PyTorch里训好的YOLO模型，变成边缘设备能认的格式。我个人习惯把这一步叫做「模型出境准备」——你得给它办齐证件，它才能在其他设备上跑起来。
3.1 PyTorch训练基础回顾
先简单过一下训练流程。YOLO的训练代码网上很多，但核心就这几步：
- 数据准备：标注好的图片，转成YOLO格式的txt标签
- 模型定义：用YOLOv5/v8/v10的官方代码，或者自己魔改
- 训练循环：前向传播 → 计算损失 → 反向传播 → 更新权重
- 验证：看mAP、Precision、Recall这些指标
嗯，这里有个坑。我在项目中遇到过，很多人训练时用的图片尺寸和部署时不一样。训练用640×640，部署时想用320×320省算力——结果模型直接崩了。为什么？因为YOLO的anchor和特征图尺寸是绑定的。
3.2 模型结构解析：YOLO到底长啥样
你想想看，YOLO本质上就是一个CNN backbone + neck + head。Backbone负责提特征，neck做特征融合，head输出检测结果。
我用YOLOv8举个例子，它的结构大致是：
| 组件 | 作用 | 常见层 |
|---|---|---|
| Backbone | 提取图像特征 | Conv + C2f + SPPF |
| Neck | 多尺度特征融合 | FPN + PAN |
| Head | 输出检测结果 | Conv + Decouple |
这里要注意的是，Head部分在部署时经常需要改造。训练时的Head包含了复杂的损失计算逻辑，但推理时只需要前向传播。我之前踩过坑，直接把训练好的模型导出，结果在边缘设备上推理速度慢得离谱——后来发现是Head里多余的reshape和concat操作没去掉。
3.3 ONNX导出：把PyTorch模型转成通用格式
ONNX（Open Neural Network Exchange）说白了就是模型界的通用语言。PyTorch模型转ONNX，用torch.onnx.export就行。但这里面的门道不少。
先看一个标准的导出代码：
import torch
import torch.onnx
# 加载训练好的模型
model = torch.load('yolov8n.pt', map_location='cpu')
model.eval()
# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 640, 640)
# 导出ONNX
torch.onnx.export(
model,
dummy_input,
'yolov8n.onnx',
input_names=['images'],
output_names=['output'],
dynamic_axes={
'images': {0: 'batch_size'},
'output': {0: 'batch_size'}
},
opset_version=12
)
这里有几个关键参数：
- opset_version：ONNX算子集的版本。我建议用12或以上，支持更多算子
- dynamic_axes：动态batch。边缘设备上一般固定batch=1，但留着这个选项没坏处
- input_names/output_names：给输入输出起个名字，后面推理时要用
3.4 ONNX简化：去掉没用的算子
导出的ONNX模型通常有很多冗余操作。比如一些恒等映射、多余的reshape、没用的concat。这些在PyTorch里可能不影响速度，但在边缘设备上就是实打实的算力浪费。
我常用的工具是onnx-simplifier：
pip install onnx-simplifier
python -m onnxsim yolov8n.onnx yolov8n_simplified.onnx \
--input-shape 1,3,640,640
简化前后对比一下：
| 指标 | 简化前 | 简化后 |
|---|---|---|
| 节点数 | 352 | 287 |
| 模型大小 | 14.2 MB | 13.8 MB |
| 推理耗时（CPU） | 45 ms | 38 ms |
你看，节点数少了将近20%，推理速度也提了15%左右。别小看这几毫秒，在边缘设备上，每一毫秒都很珍贵。
3.5 ONNX验证：确保模型没坏
导出和简化都做完了，最后一步是验证。这一步不能省，我吃过亏。
验证分三层：
- 结构验证：用onnx.checker检查模型结构是否合法
- 数值验证：用同一张图片，对比PyTorch和ONNX的输出
- 性能验证：在目标设备上跑一下，看推理时间和内存占用
数值验证的代码大概这样：
import onnx
import onnxruntime as ort
import numpy as np
# 1. 结构验证
onnx_model = onnx.load('yolov8n_simplified.onnx')
onnx.checker.check_model(onnx_model)
# 2. 数值验证
ort_session = ort.InferenceSession('yolov8n_simplified.onnx')
input_data = np.random.randn(1, 3, 640, 640).astype(np.float32)
outputs = ort_session.run(['output'], {'images': input_data})
# 和PyTorch输出对比
# 计算余弦相似度
cos_sim = np.dot(outputs[0].flatten(), torch_output.flatten()) / \
(np.linalg.norm(outputs[0].flatten()) * np.linalg.norm(torch_output.flatten()))
print(f'Cosine similarity: {cos_sim:.4f}')
3.6 本章知识体系
下面这张图总结了从PyTorch训练到ONNX导出的完整流程：
整个流程走下来，你手里的模型就从PyTorch的.pt文件，变成了一个标准化的ONNX模型。这个ONNX模型可以在任何支持ONNX Runtime的设备上运行——包括我们后面要讲的Jetson、RK3588、树莓派等等。