yolov5训练自己的数据集

---

前言

花了一些时间来用yolov5算法训练自己的数据集，其中也踩了一些坑，所以记录一下这个学习过程中的心得

---

一、参考资料

1.B站一个时长为6分钟的视频，讲解整个训练过程

2.yolov5的gitee仓库地址

3.yolov5的github仓库地址

4.labelimg的github仓库

5.labelimg的gitee仓库

我的训练过程参考了这个视频，视频中使用的工具是goole的colaboratory，需要科学上网，不过即使不能科学上网，这个视频也有很大借鉴意义，我的笔记主要针对在复现视频过程中遇到的问题，包括打标签，命令行的改变。以及之后在自己电脑，还有云主机比如恒源云上面的训练经历。

---

二、训练过程

1.首先是的打标签这个工作：

打标签具体细节可以借鉴这篇博客——点击即可跳转访问

①打标签的软件是：labelImg 。

②打标签时的命令行启动为：
pyrcc5 -o libs/resources.py resources.qrc
python labelImg.py
其在第一次cmd进去时，记得导入需求包。

③打标签前可以修改其原文件夹下的data目录下的predefined_classes。 将其设置为自己所分的几个类。如果不提前修改，nc 即 names count需要加上软件里面预设的names数量。

④同时在需要打标签的图片文件夹里面，图片集images的同级目录下面需要有自己的创建的labels（标签存放地址），以及 .yaml文件， .yaml文件里面包含训练集，测试集图片存放地址，同时包含nc即标签类别总数，还有各个标签的名字，需注意， .yaml文件对格式要求严格，nc： ，names： 等后面必须先是一个空格才是输入内容。（labels和yaml文件是自己创建的，不是自动生成的，labels文件夹一开始是空的文件夹，随着打标签之后标签的保存，里面会生成class.xxx和标签，生成的文件类别由打标签软件功能的选择有关，yaml文件里面的内容是自己填写的，示例如下图）：

如果不进行上述工作，可能报类似以下错误：

AssertionError: Label class 16 exceeds nc=2 in /content/sleep/sleep.yaml. Possible class labels are 0-1

还有对.yaml文件里面的格式进行报错。

⑤以上准备工作完成，才可开始进行打标签工作，同时打标签时需注意你所需要的对应格式，是yolo的txt，还是xml等。一定要注意再保存。

2.然后是训练这个工作：

①先下载yolov5这个包，根据网络选择gitee或者github上面的仓库导入。然后进入yolov5文件夹里面，安装其需求文档，安装完成之后检查安装的torch版本是否为GPU版本，如果不是，需要卸载torch和torchvision之后去pytorch官网找到命令行重装torch的GPU版本，之后再行检查。
检查代码：

import torch 
print(torch.__version__) 0.4.0    #显示torch版本
print(torch.cuda.is_available())   #显示是否可用gpu，ture为可用，flase为不可用，需要重新装gpu版本的。
1
2
3

②如果是在goole的colaboratory上面训练的话，需要导入训练集的压缩包，导入过程稍微有点慢，需要耐心等待，在colaboratory里面因为很多时候会不清楚自己所处路径问题，所以路径都建议使用绝对路径。

③修改

yolov5/models/yolov5.yaml

文件里面的nc与训练集的nc保持一致。（虽然不改也能运行，但是保不齐出什么问题，比如训练出来的权重文件不完美）

④开始训练：--data /content/mytest/rubbish.yaml根据自身文件存放位置决定，--img 640 --batch 64 是一次传入的图片数量及训练次数，根据选择的训练模型不同，GPU显存的占用情况也不同，以使用模型为yolov5s，img为640✖️640为例，goole的colaboratory上面一次可以传入64张图片，恒源云上面显存为24G的RTX3090可以一次传入128张，在图片尺寸不变的情况下，如果训练时选用模型为yolov5l，因为结构更为复杂，所以一次可以传入训练的图片数量更少。训练次数–batch一般需要较多，否则可能连基本的识别框都训练不出来。

在goole的colaboratory上命令为:(colaboratory有一块显存为12G的Tesla K80显卡)

!python /content/yolov5/train.py --img 640 --batch 64 --epochs 300 --data /content/mytest/rubbish.yaml --cfg yolov5s.yaml --weights yolov5s.pt --nosave --cache --device '0'）
1

不加上–device '0’可能会报以下错误：

warnings.warn(‘User provided device_type of ‘cuda’, but CUDA is not available. Disabling’）

在自己的电脑上面训练命令为：

python /content/yolov5/train.py --img 640 --batch 64 --epochs 300 --data /content/mytest/rubbish.yaml --cfg yolov5s.yaml --weights yolov5s.pt --nosave --cache
1

⑤训练结束后得到新的权重文件：正常情况下会存放在

yolov5/runs/train

的最新的exp文件夹下面的weights文件夹里面的 .pt文件，一般训练结束之后都会有提示。如不是在本机训练，记得保存这个权重文件。

⑥测试训练效果命令行：（分别逐条运行）

from IPython.display import Image
Image(filename="/content/yolov5/runs/train/exp/train_batch0.jpg",width=800)
	
Image(filename="/content/yolov5/runs/train/exp/train_batch1.jpg",width=800)
	
Image(filename="/content/yolov5/runs/train/exp/train_batch2.jpg",width=800)
	
Image(filename="/content/yolov5/runs/train/exp/labels_correlogram.jpg",width=800)
	
Image(filename="/content/yolov5/runs/train/exp/results.png",width=800)
1
2
3
4
5
6
7
8
9
10

会依次出现图像测试图和统计图形。

⑦运行yolov5算法，进行测试：将新的权重文件放到

yolov5/weight

里面，然后修改

yolov5/detect.py

里面的代码

parser.add_argument('--weights', nargs='+', type=str, default='weights/last.pt', help='model path(s)')
1

，在 default='weights/last.pt'里放置你的权重文件路径，开始运行yolov5算法运行即可：（修改权重的语句一般放在一大堆parser.add_argument文件的第一行，随着版本迭代，可能具体代码有所差别，但是其实也都大同小异）

python detect.py --source 0  # webcam
                          img.jpg  # image 
                          vid.mp4  # video
                          path/  # directory
                          path/*.jpg  # glob
                          'https://youtu.be/Zgi9g1ksQHc'  # YouTube
                          'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
1
2
3
4
5
6
7

关于yolov5算法复现具体过程，可以参考其他资料。

总结

yolov5很神奇，训练自己的训练集也很有趣。