实践教程｜如何用YOLOv7训练自己的数据集

作者丨Hernando

编辑丨极市平台

导    读

 

使用YOLOv7训练数据集的全过程详解。

一、设备准备

我用的是V100*2，总共64GB显卡内存，在Ubuntu 18.04系统，python 3.7.10下运行。

二、代码克隆，环境安装与权重下载

1. 代码克隆

代码的克隆和环境的安装在官网上也有指导，可以直接跳转到官网 (https://github.com/WongKinYiu/yolov7) 查看。下面是具体的操作，跟官网的一样。首先，直接把项目从官网上克隆下来，在终端输入如下命令

git clone https://github.com/WongKinYiu/yolov7

代码克隆过程

代码的克隆过程看服务器的网速，有的需要两三分钟，有的几十秒即可。克隆好了之后，打开yolov7，可以看到如下的文件。

2. 环境安装

然后进入到项目文件夹中，进行环境的安装，分别输入一下两行代码即可。同样地，速度受服务器网速影响，但是都不会太久。

cd yolov7
pip install -r requirements.txt

下面是环境安装好之后的结果截图。

环境安装好之后

3. 权重下载

这次我们使用的预训练模型是​​yolov7_training.pt​​​，这个模型是在MS COCO数据集上进行训练的，我们把这个模型下载到​​/yolov7/weights/​​​下。注意，我们需要创建一个​​/weights/​​文件夹，用来统一存放我们的预训练模型。输入如下命令进行文件夹的创建，并进入到该文件夹下

mkdir weights   # 创建文件夹
cd weights      # 进入到该文件夹下

这些预训练的模型官网都有提供，掌握如何训练自己的数据集后可根据需求进行预训练模型的更改。官网的截图如下：

官网提供的预训练模型

在终端输入如下命令，即可下载预训练模型

wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7_training.pt

下图是下载过程的截图，同样是受服务器的网速影响。可以看到，我只有十几KB/s，而整个文件大小为75.6MB，所以这个速度是不够的。这里建议手动下载，然后手动传到服务器上，这样速度更快。

下载完后（或者上传好后，我是手动下载再上传的），会在​​/weights/​​得到一个如下图的文件

然后我们回到上一级工作路径下，也就是​​/yolov7/​​下，输入如下命令

cd ..

到这里，我们的预训练权重也就下载好了。

4. 测试一下

这一步的重要目的，是测试环境是否都安装好了，如果能够顺利跑下来并得到相应的结果，就说明环境是安装好的了；如果没安装好，则再安装一次环境即可，也就是第2步。

首先，我们下载已经训练好的权重文件（这里的权重文件和Transfer Learning的预训练文件并不是同一个文件），同样地，下载到​​/weights/​​下面。

cd weights
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt
cd ..

下载好了之后，得到两个权重文件。接下来，就是直接进行测试了。现在我们在​​/yolov7/​​这个工作路径下，输入如下命令

python detect.py --weights weights/yolov7.pt --source inference/images

参数说明

--weights weight/yolov7.pt   # 这个参数是把已经训练好的模型路径传进去，就是刚刚下载的文件
--source inference/images   # 传进去要预测的图片

如果得到如下的运行结果，则说明运行成功，预测的图片被保存在了​​/runs/detect/exp/​​文件夹下

查看其中一张图片，如下

推理前，在source里

推理后

到此，就说明环境安装好了，可以进行下面的数据准备工作。

三、数据准备

1. 准备工作

我们生成​​/datasets/​​​文件夹，把数据都放进这个文件夹里进行统一管理。训练数据用的是yolo数据格式，不过多了两个​​.txt​​文件，这两个文件存放的，是每个图片的路径，后面会具体介绍。

2. 准备yolo格式的数据

2.1. 创建文件夹

mkdir datasets

得到如下的结果

2.2. 准备yolo格式数据集

如果不懂yolo格式数据集是什么样子的，建议先学习一下。大部分CVer都会推荐用labelImg进行数据的标注，我也不例外，推荐大家用labelImg进行数据标注。不过这里我不再详细介绍如何使用labelImg，网上有很多的教程。同时，标注数据需要用到图形交互界面，远程服务器就不太方便了，因此建议在本地电脑上标注好后再上传到服务器上。

这里假设我们已经得到标注好的yolo格式数据集，那么这个数据集将会按照如下的格式进行存放。

不过在这里面，​​train_list.txt​​​和​​val_list.txt​​是后来我们要自己生成的，而不是labelImg生成的；其他的则是labelImg生成的。

接下来，就是生成 ​​train_list.txt​​​和​​val_list.txt​​​。​​train_list.txt​​​存放了所有训练图片的路径，​​val_list.txt​​则是存放了所有验证图片的路径，如下图所示，一行代表一个图片的路径。这两个文件的生成写个循环就可以了，不算难。

train_list.txt内容示例

到此，数据集就已经准备好了。

四、配置训练的相关文件

1. 前言

总共有两个文件需要配置，一个是​​/yolov7/cfg/training/yolov7.yaml​​​，这个文件是有关模型的配置文件；一个是​​/yolov7/data/coco.yaml​​，这个是数据集的配置文件。

2. 配置模型文件

第一步，复制​​yolov7.yaml​​​文件到相同的路径下，然后重命名，我们重命名为​​yolov7-Helmet.yaml​​。

第二步，打开​​yolov7-Helmet.yaml​​文件，进行如下图所示的修改，这里修改的地方只有一处，就是把nc修改为我们数据集的目标总数即可。然后保存。

模型配置文件修改展示

3. 配置数据集文件

第一步，复制​​coco.yaml​​​文件到相同的路径下，然后重命名，我们命名为​​Helmet.yaml​​。

第二步，打开​​Helmet.yaml​​​文件，进行如下所示的修改，需要修改的地方为5处。第一处：把代码自动下载COCO数据集的命令注释掉，以防代码自动下载数据集占用内存；第二处：修改train的位置为​​train_list.txt​​​的路径；第三处：修改val的位置为​​val_list.txt​​的路径；第四处：修改nc为数据集目标总数；第五处：修改names为数据集所有目标的名称。然后保存。

数据配置文件修改展示

至此，相关的文件已经配置完成。

五、开始训练

1. 前言

训练用到之前提到的三个文件：预训练模型​​yolov7_training.pt​​​、​​yolov7-Helmet.yaml​​​和​​Helmet.yaml​​，当然还有一些其他的超参数，具体的命令如下：

python train.py --weights weights/yolov7.pt --cfg cfg/training/yolov7-Helmet.yaml --data data/Helmet.yaml --device 0,1 --batch-size 64 --epoch 10

然后运行界面如下，图中红色部分表示开始扫描数据，然后会在​​yolov7/datasets/Helmet/​​下生成两个缓存文件，这两个文件是方便模型进行数据读取生成的，训练过程中不要删除。扫描完后模型会开始训练，如下图。

模型训练中

下面简单介绍一下相关的参数。

--weights weights/yolov7.pt             # 接收预训练模型路径的参数
--cfg cfg/training/yolov7-Helmet.yaml   # 接收模型配置文件的参数
--data data/Helmet.yaml                 # 接收数据配置文件的参数
--device 0,1                            # GPU/CPU训练，我有2块，因此为0,1；若1块，则0；若CPU，则cpu
--batch-size 64                         # 按照自己GPU内存大小大致确定
--epoch 10                              # 不用多说；我仅展示，因此不训练那么多次了

更多的参数可以从​​train.py​​文件中找到，如下图

train.py参数

六、训练结果与推理

1. 训练结果

训练结束后，终端会打印出最好的模型和最后一个epoch的模型结果保存在哪里，如下图所示，

训练结束

在​​/runs/train/exp/​​​下，也存了其他的训练结果文件，比如每个epoch的结果到保存到了​​.txt​​文件夹里，如下图所示；有兴趣的可以自行探索。

2. 推理

推理其实就是跟​​二-4. 测试一下​​里的差不多，无非就是把相关的权重文件和测试图片换一下即可。

python detect.py --weights runs/train/exp3/weights/best.pt --source datasets/testImages

推理的结果如下所示

七、总结

我也是个小白，这篇文章记录了我对yolov7的使用过程，仅作交流分享，如有侵权请联系我。如有问题，欢迎指出。

本文仅做学术分享，如有侵权，请联系删文。​

—THE END—