如何基于PyTorch/TensorFlow用自定义数据集训练YOLOv3/v4实现实时目标检测

全程不用碰Darknet的C代码，全Python流程，你已经有YOLO格式标注的话不用做格式转换，按下面步骤走就能复现，我自己在工业检测场景跑过不下10次，没有虚的步骤。

1. 环境准备 #

• 先建个独立Python虚拟环境，版本选3.8~3.10，兼容性最好，不会出各种奇奇怪怪的依赖冲突
• 先装匹配自己显卡CUDA版本的PyTorch，有N卡一定要装CUDA版，训练速度比CPU快几十倍；之后补装opencv-python、pyyaml、tqdm、matplotlib这几个基础依赖就行，不用装多余的包
• 拉取官方维护的PyTorch版YOLO代码到本地，根目录直接操作，不需要编译，开箱能用

2. 数据集规整（你已有YOLO格式标注，这步10分钟就能弄完） #

按下面的结构放文件就行，标注内容不用改：

your_dataset/
├── images/
│   ├── train/  # 放训练集图片，jpg、png格式都支持
│   └── val/    # 放验证集图片，全量数据按8:2或者7:3拆分就行，不用特意抠比例
└── labels/
    ├── train/  # 放和训练集图片一一对应的txt标注，比如test.jpg对应test.txt
    └── val/    # 放和验证集图片一一对应的txt标注

注意核对标注格式：txt里每一行是class_id x_center y_center bbox_width bbox_height，所有坐标必须是归一化到0~1区间的数值；你是9类目标，class_id就从0到8连续编号，别跳号、别超范围，不然训练必报错。

3. 配置文件修改 #

总共就改两个配置，不用动核心源码：

• 新建数据集配置文件custom9.yaml，写清楚三个核心参数就行：

  path: ./your_dataset  # 填你自己数据集的根目录，相对路径绝对路径都可以
  train: images/train
  val: images/val
  nc: 9 # 固定写9，对应你的9个检测类别
  names: ['类别0名称','类别1名称','类别2名称','类别3名称','类别4名称','类别5名称','类别6名称','类别7名称','类别8名称'] # 按class_id顺序填你自己的类别名
  

• 复制对应版本的模型配置文件（用YOLOv3就复制yolov3.yaml，用v4就复制yolov4.yaml），把文件里两处nc参数的默认值80改成9就行，其他网络结构参数别动，直接适配预训练权重的结构。

4. 预训练权重准备 #

下载对应版本的PyTorch格式预训练权重，放到项目的weights目录下就行，不用从头初始化训练，迁移学习收敛速度快3~5倍，小数据集也能出不错的精度。

5. 启动训练 #

直接在项目根目录跑训练脚本，常用参数给你列好，按需调整：

python train.py --weights ./weights/yolov3.pt --cfg ./models/yolov3-custom.yaml --data ./data/custom9.yaml --epochs 100 --batch-size 16 --img 640 --device 0

参数说明：

• --device 0代表用第一块N卡训练，没有显卡就填cpu，就是速度慢很多
• --batch-size按自己显卡显存调，显存不够就往小了改（比如8、4都可以），只影响训练速度，不影响最终精度
• --img 640是训练输入分辨率，实时检测场景640足够用，要更高精度可以改1280，但显存占用会涨很多
• 训练过程中会自动在runs目录下保存每轮权重、验证集精度、混淆矩阵、PR曲线，不用自己额外写代码算指标，盯着验证集的mAP指标就行，取mAP最高的best.pt做最终部署权重。

6. 推理&实时检测验证 #

训练完直接拿best权重跑推理就行，实时检测直接调摄像头：

python detect.py --weights ./runs/train/exp/weights/best.pt --source 0 --conf-thres 0.25 --iou-thres 0.45

• --source 0就是调用本地默认摄像头，也可以传视频路径、图片文件夹路径做批量测试
• 速度方面，单张1080Ti跑640分辨率的YOLOv3能到100FPS以上，用轻量化的tiny版本能到200FPS以上，完全满足实时检测要求。

常见踩坑避坑 #

• 一定要先核对图片和标注的文件名一一对应，别出现有图没标注、有标注没图的情况，花2分钟写个10行以内的Python脚本扫一遍文件名匹配度就行，省得训练中途报错
• 训练前先抽10张图跑标注可视化，确认框的位置、类别和实际目标对得上，别标错了训好几天白忙活
• 如果训练时loss一直不下降，第一时间检查标注坐标是不是归一化的，很多人导出标注时没开归一化，存的是原始像素坐标，肯定训不收敛
• 别乱改网络结构里的anchor、通道数这些默认参数，用预训练权重默认的参数就行，自定义数据量不大的话改了反而容易掉点

内容的提问来源于stack exchange，提问作者Ecu