从0开始学视觉（5）——利用yolov5训练自己的目标检测模型

目录

1.下载源码

2. 环境的配置和依赖库的安装

3. 数据集和预训练权重的准备 

4.训练自己的模型 

5. 启用tensorbord查看参数

 6.测试集

---

1.下载源码

        首先打开yolov5的github的官网https://github.com/ultralytics/yolov5/tree/v5.0（这个网站在国外打开是很慢的，而且是有的时候能正常打开，有的时候是进不去的，但是大家第一次打不开的话，一定要多打开几次。）打开的官网界面如下，这个就是大神glenn-jocher开源的yolov5的项目。

        将我们下载好的yolov5的代码解压，然后用一款IDE打开（我用的是pycharm），打开之后整个代码目录如下图：

 

         现在来对代码的整体目录做一个介绍：

        ├── data：主要是存放一些超参数的配置文件（这些文件（yaml文件）是用来配置训练集和测试集还有验证集的路径的，其中还包括目标检测的种类数和种类的名称）；还有一些官方提供测试的图片。如果是训练自己的数据集的话，那么就需要修改其中的yaml文件。但是自己的数据集不建议放在这个路径下面，而是建议把数据集放到yolov5项目的同级目录下面。

        ├── models：里面主要是一些网络构建的配置文件和函数，其中包含了该项目的四个不同的版本，分别为是s、m、l、x。从名字就可以看出，这几个版本的大小。他们的检测测度分别都是从快到慢，但是精确度分别是从低到高。这就是所谓的鱼和熊掌不可兼得。如果训练自己的数据集的话，就需要修改这里面相对应的yaml文件来训练自己模型。

        ├── utils：存放的是工具类的函数，里面有loss函数，metrics函数，plots函数等等。

        ├── weights：放置训练好的权重参数。

        ├── detect.py：利用训练好的权重参数进行目标检测，可以进行图像、视频和摄像头的检测。

        ├── train.py：训练自己的数据集的函数。

        ├── test.py：测试训练的结果的函数。

        ├──requirements.txt：这是一个文本文件，里面写着使用yolov5项目的环境依赖包的一些版本，可以利用该文本导入相应版本的包。

        以上就是yolov5项目代码的整体介绍。我们训练和测试自己的数据集基本就是利用到如上的代码。

2. 环境的配置和依赖库的安装

        打开requirements.txt这个文件，可以看到里面有很多的依赖库和其对应的版本要求。我们打开pycharm的命令终端，在中输入如下的命令，就可以安装了。

pip install -r requirements.txt

3. 数据集和预训练权重的准备 

        把前面我们已经标注好的数据集放进来，数据集的目录格式如下图所示。

        预训练权重可以通过这个网址（https://github.com/ultralytics/yolov5/releases）进行下载，本次训练自己的数据集用的预训练权重为yolov5s.pt。 

4.训练自己的模型 

         预训练模型和数据集都准备好了，就可以开始训练自己的yolov5目标检测模型了，训练目标检测模型需要修改两个yaml文件中的参数。一个是data目录下的相应的yaml文件，一个是model目录文件下的相应的yaml文件。

        修改data目录下的相应的yaml文件。找到目录下的voc.yaml文件，将该文件复制一份，将复制的文件重命名，最好和项目相关，这样方便后面操作。我这里修改为hat.yaml。该项目是对安全帽的识别。

        打开这个文件夹修改其中的参数，首先将箭头1中的那一行代码注释掉（我已经注释掉了），如果不注释这行代码训练的时候会报错；箭头2中需要将训练和测试的数据集的路径填上（最好要填绝对路径，有时候由目录结构的问题会莫名奇妙的报错）；箭头3中需要检测的类别数，我这里是识别安全帽和人，所以这里填写2；最后箭头4中填写需要识别的类别的名字（必须是英文，否则会乱码识别不出来）。到这里和data目录下的yaml文件就修改好了。

 

        由于该项目使用的是yolov5s.pt这个预训练权重，所以要使用models目录下的yolov5s.yaml文件中的相应参数（因为不同的预训练权重对应着不同的网络层数，所以用错预训练权重会报错）。同上修改data目录下的yaml文件一样，我们最好将yolov5s.yaml文件复制一份，然后将其重命名，我将其重命名为yolov5_hat.yaml。打开yolov5_hat.yaml文件只需要修改如图中的数字就好了，这里是识别两个类别。

        找到train.py这个py文件。找到主函数的入口，这里面有模型的主要参数，训练自己的模型需要修改如下几个参数就可以训练了。首先将weights权重的路径填写到对应的参数里面，然后将修好好的models模型的yolov5s.yaml文件路径填写到相应的参数里面，最后将data数据的hat.yaml文件路径填写到相对于的参数里面。这几个参数就必须要修改的参数。 

还有几个需要根据自己的需求来更改的参数：首先是模型的训练轮次，这里是训练的30轮。 

        其次是输入图片的数量和工作的核心数，这里每个人的电脑都不一样，所以这里每个人和自己的电脑的性能来。这里可以根据我的电脑的配置做参考，我的电脑是拯救者2060版本的显卡，cpu的核心数是16核。我的电脑按默认的参数输入图片数量为16，工作核心为8的话就会出现GPU显存溢出的报错。每个人的电脑配置不一样，所以可以根据自己的电脑配置来修改参数。

但是pycharm的用户可能会出现如下的报错。这是说明虚拟内存不够了。 

可以在utils路径下找到datasets.py这个文件，将里面的第81行里面的参数nw改完0就可以了。运行train.py函数训练自己的模型了。 

5. 启用tensorbord查看参数

         yolov5里面有写好的tensorbord函数，可以运行命令就可以调用tensorbord，然后查看tensorbord了。首先打开pycharm的命令控制终端，输入：

tensorboard --logdir=runs/train

就会出现一个网址地址，将那行网址复制下来到浏览器打开就可以看到训练的过程了

如果模型已经训练好了，但是我们还想用tensorbord查看此模型的训练过程，就需要输入如下的命令。就可以看到模型的训练结果了。 

tensorboard --logdir=runs

 6.测试集

        等到数据训练好了以后，就会在主目录下产生一个run文件夹，在run/train/exp/weights目录下会产生两个权重文件，一个是最后一轮的权重文件，一个是最好的权重文件，一会我们就要利用这个最好的权重文件来做测试。除此以外还会产生一些验证文件的图片等一些文件。找到主目录下的detect.py文件，打开该文件。然后找到主函数的入口，这里面有模型的主要参数。

这里需要将刚刚训练好的最好的权重传入到推理函数中去。然后就可以对图像视频进行推理了。

 

对图片进行测试推理，将如下参数修改成图片的路径，然后运行detect.py就可以进行测试了。 

推理测试结束以后，在run下面会生成一个detect目录，推理结果会保存在exp目录下。

 

对视频进行测试，和如上的图片的测试是一样的，只不过是将图片的路径改为视频的路径而已。

 

 利用摄像头进行测试只需将路径改写为0就好了。但是11报错如下。

解决方法：首先找到datasets.py这个py文件。  打开文件，找到第279行代码，给两个url参数加上str就可以了，如图所示，就可以完美运行电脑的摄像头了。