【图文并茂】【一篇搞定】yolov5训练自己的数据集（交通训练集）

本文总项目文件：yolov5-6.0

---

😝1 安装库😝

💘💘1.1 创建虚拟环境💘💘

创建python3.7环境

conda create -n yolov5 python=3.7

💘💘1.2 下载官网的yolov5💘💘

去官网 ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)
这里下载yolov5的v6.0版本，而最新是v7.0，可能需要python3.8以上才行。

💘💘1.3 安装依赖💘💘

进入yolov5-6.0目录 输入cmd，激活环境，方便使用该目录下的requirements安装相关依赖

修改为临时pip清华源下载：

pip config set global.index-url --site https://pypi.tuna.tsinghua.edu.cn/simple

根据requirements安装依赖

pip install -r requirements.txt

注：pip install -r requirement.txt 默认安装torch的cpu版本！

pycharm打开目录文件夹，并配置conda解释器

---

---

---

😝2 准备图片数据😝

本文所采用图片：开源数据集：交通标志数据集CCTSDB2021

💘💘2.1 给图片打label💘💘

安装labelimg

pip install labelimg

命令行里启动labelimg

打开文件夹

改为yolo格式

打标签

点击左侧save保存为txt格式

💘💘2.2 放置图片💘💘

yolov5 - 6.0 目录下新建dataset文件夹

dataset
- images  （存放原图）
	- train
		- 1.jpg
		- 2.jpg
		- ....
	- val
		- 101.jpg
		- 102.jpg
		- ....
- labels  （存放刚刚用labelimg打完标签的txt文件）
	- train
		- 1.txt
		- 2.txt
		- ....
	- val
		- 101.txt
		- 102.txt
		- ....

同理dataset/labels也是一样，只不过存放的是刚刚用labelimg打完标签的txt文件，文件名称要一一对应好

💘💘2.3 建立数据配置💘💘

yolov5 - 6.0 / data 目录下新建traffic_logo.yaml
模板内容如下：

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, >path/to/imgs2, ..]  
path: D:\Desktop\yolov5-6.0\dataset  # dataset root dir  
train: images/train  # train images (relative to 'path')  
val: images/val  # val images (relative to 'path')  
test:  # test images (optional)  

# Classes  
nc: 3  # number of classes  
names: ['prohibitory','warning','mandatory']  # class names
- path：数据集的根目录
- train：训练集与path的相对路径
- val：验证集与path的相对路径
- nc：类别数量，因为这个数据集只有3个类别（fire），nc即为3。
- names：类别名字。

---

---

---

😝3 训练😝

💘💘3.1 下载预训练权重💘💘

下载coco的预训练模型权重 yolov5s.pt 文件：
Release v6.0 - YOLOv5n ‘Nano’ models, Roboflow integration, TensorFlow export, OpenCV DNN support · ultralytics/yolov5 (github.com)

放置在yolov5-6.0 根目录下

💘💘3.2 开始训练💘💘

pycharm在左下角的终端打开命令行

输入训练命令：

python train.py --data data/mydata_traffic.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt -->epoch 10 --batch-size 2 --device cpu

• --data：数据配置
• --cfg：模型配置
• --weights：这里采用已经经过coco数据集训练过的预训练模型（迁移学习）
• --device：gpu则写0，cpu则写cpu
  

不出意外地会出小意外：

Traceback (most recent call last):
  File "train.py", line 620, in <module>
    main(opt)
  File "train.py", line 517, in main
    train(opt.hyp, opt, device, callbacks)
  File "train.py", line 316, in train
    loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
  File "D:\Desktop\yolov5-6.0\utils\loss.py", line 120, in __call__
    tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets
  File "D:\Desktop\yolov5-6.0\utils\loss.py", line 217, in build_targets
    indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
RuntimeError: result type Float can't be cast to the desired output type __int64

根据提示，找到对应文件修改即可：点击下方橙色处 跳转 到这个文件

修改底173处：原语句后面增加.long()即可

然后再次运行！成功训练！

训练完成所有结果保存在run/train/exp5中.

---

---

---

😝4 测试😝

输入测试命令：

python detect.py --weights runs/train/exp5/weights/best.pt --source  dataset/images/val/00401.jpg

• –weights：使用权重，这里测试就使用我们刚刚训练出来的best.pt即可
• –source: 测试图片的路径

运行结果保存在run/detect/exp中

查看预测图00401.jpg
（由于硬件和时间原因，我只训练了10个epoch，相信充足的训练效果更好！）