2、YoloV5模型训练与导出：从源码到ONNX格式

好，咱们进入正题。这一章，我带你走一遍YoloV5的完整流程——从源码结构、数据集准备，到模型训练、评估，最后导出为ONNX格式。嗯，每一步我都会结合自己的踩坑经验来讲，希望能帮你少走弯路。

2.1 YoloV5源码结构

YoloV5的源码，说白了就是一套基于PyTorch的目标检测框架。我第一次看源码时，说实话有点懵——文件太多了。但后来我发现，核心文件就那么几个。

我个人习惯把源码分成三块来看：

  • 模型定义：models/ 目录下的 yolo.py、common.py。这里定义了YoloV5的网络结构，包括Backbone、Neck、Head。
  • 训练与推理：train.py 和 detect.py。一个负责训练，一个负责推理。
  • 工具函数：utils/ 目录。包括数据加载、损失计算、评估指标等。

你可能会问：“这么多文件，我该从哪看起？”我的建议是——先跑通 train.py，再回头读代码。我在项目中遇到过很多次，光看代码不跑，根本理解不了参数的作用。

小提示：YoloV5的源码结构其实很清晰。你只要记住：models 管网络，utils 管工具，data 管数据。其他都是辅助。

2.2 数据集准备

数据集准备，说白了就是让你的数据能被YoloV5读懂。YoloV5支持两种格式：COCO格式和YOLO格式。我个人更推荐YOLO格式，因为它简单直接。

YOLO格式的数据集长这样：

dataset/
├── images/
│   ├── train/
│   │   ├── img1.jpg
│   │   └── img2.jpg
│   └── val/
│       ├── img3.jpg
│       └── img4.jpg
└── labels/
    ├── train/
    │   ├── img1.txt
    │   └── img2.txt
    └── val/
        ├── img3.txt
        └── img4.txt

每个 .txt 文件对应一张图片的标注。格式是：

class_id x_center y_center width height

注意，这里的坐标都是归一化后的，范围在0到1之间。我曾经犯过一个低级错误——直接用像素坐标写进去，结果训练出来的模型完全没法用。嗯，这个坑你千万别踩。

另外，你还需要一个 data.yaml 文件，告诉YoloV5你的数据集在哪：

train: dataset/images/train
val: dataset/images/val

nc: 2  # 类别数量
names: ['cat', 'dog']  # 类别名称
注意：YoloV5对数据集路径很敏感。我建议用绝对路径，或者把数据集放在YoloV5源码目录下。否则训练时会报“找不到图片”的错误。

2.3 模型训练

训练命令其实很简单，一行代码搞定：

python train.py --data data.yaml --cfg models/yolov5s.yaml --weights '' --batch-size 16 --epochs 100

这里我解释几个关键参数：

  • --data：指向你的 data.yaml 文件。
  • --cfg：模型配置文件。YoloV5提供了 yolov5s、yolov5m、yolov5l、yolov5x 四种规格。s最小最快，x最大最准。
  • --weights：预训练权重。如果为空，则从头训练。
  • --batch-size：批次大小。根据你的显存来调。
  • --epochs：训练轮数。

我个人习惯先用 yolov5s 跑一轮，看看数据集有没有问题。没问题了再换大模型。我在项目中遇到过，有人一上来就用 yolov5x，结果显存爆了，还找不到原因。

训练过程中，你会看到类似这样的输出：

     Epoch   gpu_mem      box       obj       cls     labels   img_size
     1/100     2.45G   0.0896   0.0215   0.0000        10       640

这里 box、obj、cls 是三个损失值。它们都在下降，说明模型在学。如果 obj 一直降不下去，嗯，可能是数据标注有问题。

小技巧：训练时建议开启 --exist-ok 参数，这样如果训练中断了，可以接着上次的权重继续训练，不用从头开始。

2.4 模型评估

训练完成后，YoloV5会自动在验证集上评估。但你也可以手动跑评估：

python val.py --data data.yaml --weights runs/train/exp/weights/best.pt

评估结果会输出mAP（平均精度均值）。一般来说，mAP@0.5 达到0.9以上就算不错了。但注意，mAP高不代表模型在实际场景中好用。我曾经遇到过，mAP 0.95的模型，在真实摄像头下表现一塌糊涂——因为训练集和实际场景的光照条件差太多了。

评估结果还会生成一些图表，比如混淆矩阵、PR曲线。这些图能帮你分析模型哪里做得好，哪里做得不好。我个人习惯先看混淆矩阵——如果某个类别经常被误检成另一个类别，那可能是这两个类别太像了，需要增加训练数据。

2.5 导出为ONNX格式

终于到了最关键的一步——导出ONNX。YoloV5官方提供了导出脚本：

python export.py --weights runs/train/exp/weights/best.pt --include onnx

就这么简单？嗯，但有几个坑要注意。

首先，导出时默认的输入尺寸是640x640。如果你的模型训练时用了其他尺寸，记得加上 --img-size 参数：

python export.py --weights best.pt --include onnx --img-size 416 416

其次，导出后的ONNX模型，建议用 onnx-simplifier 简化一下：

python -m onnxsim best.onnx best_simplified.onnx

为什么要简化？因为YoloV5的导出过程会生成一些冗余的算子，比如 Reshape、Transpose 等。这些算子在不同推理引擎上可能不支持。简化后，模型更干净，兼容性更好。

最后，验证一下导出的ONNX模型是否正确：

python -m onnxruntime.quantization.preprocess --input best.onnx --output best_verified.onnx

如果没报错，恭喜你，ONNX模型导出成功了。

重要提醒：ONNX模型只是中间格式，不是最终部署格式。你还需要根据目标平台（比如TensorRT、OpenVINO、ONNX Runtime）做进一步优化。这部分内容，我会在后面的章节详细讲。

好了，这一章的内容就到这里。从源码结构到数据集准备，从模型训练到评估，再到ONNX导出，每一步我都尽量把关键点和坑点讲清楚。下一章，我会带你深入ONNX模型的结构，看看它到底长什么样。

← 上一章 📖 返回目录 下一章 →