YOLOV5 训练自己的数据集

YOLOV5是目标检测领域，one stage类型网络中的成熟算法。本文将针对一个Finger识别项目，介绍ubuntu命令行下，yolov5从环境搭建到模型训练的整个过程。

由于需要自行制作数据集，因此标注工具直接使用yolo开发者提供的标注工具yolomark，避免在数据转换上花费过多精力。

算法原理阅读:
GitHub
深入浅出Yolo系列之Yolov5核心基础知识完整讲解
Roboflow的blog

1.环境

训练使用的环境如下：
Ubuntu 20.04
python 3.8.8
torch 1.8.2+cu111
torchvision 0.9.2+cu111
NVIDIA driver 470.74
CUDA 11.5

Ubuntu20 安装教程

Torch安装教程

Ubuntu环境安装显卡驱动教程

CUDA&CUDNN安装教程

查看Ubuntu版本号

cat /proc/version

查看Python版本号

python --version

查看torch等版本号

pip3 list

查看显卡驱动版本

nvidia-smi

查看CUDA版本

nvcc --version

2.代码模型下载

2.1 下载模型

git clone https://github.com/ultralytics/yolov5

2.2下载预训练模型

链接: https://pan.baidu.com/s/1L5t3Uis0w1nNDyEboKg5Hg
密码: 74bu

2.3 下载模型依赖项

cd yolov5
pip3 install -r requirements.txt

如果报错，打开根目录下的requirements.txt，手动安装

sudo apt install xxx

matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
Pillow>=7.1.2
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0

3.数据集制作

3.1 下载Yolomark
使用Yolomark的标注工具。

git clone https://github.com/AlexeyAB/Yolo_mark

3.2 数据转移
将图片放置到 Yolomark/x64/Release/data/img 文件夹下，删除其中预带的图片。

3.3 编译和运行Yolomark

用cmake编译.sh文件，然后赋予yolo_mark运行权限，执行

cd Yolo_mark
cmake .
make
chomd +x ./linux_mark.sh
./linux_mark.sh

3.4 标注配置
在标注之前先指定标注类型的数量和名称。
修改 Yolo_mark/x64/Release/data/obj.data , 根据实际情况指定标注类型数量。此例子只有1个类型。

classes= 1

修改 Yolo_mark/x64/Release/data/obj.names，根据实际情况指定标注类型名称。

Finger

3.5 标注
主界面如下。
框选 - 用来标注，
object ID - 选择切换标准类型
快捷键C - 清除当前图像标记
快捷键空格 - 保存并翻到下一张图片
快捷键Esc - 退出

3.6 结果生成
标注完成后，自动生成标签的txt，例如：

0 0.612500 0.174306 0.071875 0.184722
0 0.683984 0.120833 0.071094 0.183333
0 0.756641 0.097917 0.077344 0.168056
0 0.833594 0.134028 0.065625 0.145833
0 0.880859 0.311111 0.064844 0.155556

4. 训练设置

4.1 文件路径设置
yolov5的根目录如下图片，建议按照如下放置各文件
data — 放数据集和训练所用.yaml文件
models — 放yolo5模型
runs — 训练结果
weights — 放预训练模型权重
根目录 — 放train.py (训练执行入口)

4.2数据集设置
在yolov5/data下新建一个文件夹FingerDatasets，其中建立两个文件夹images和labels，放入放置标注后的jpg图片和txt标签。

4.3 训练集测试集分配
需要将标注图片分为train / test / val三种用途。可以新建一个py文件，用random函数将它们按照一定比例分配成三组。

当然，推荐手动写txt 。引导文件格式简单，Ctrl+C/V/H比码代码快多了 😃

最终生成三个txt文件：

txt内容如下例

data/FingerDatasets/images/71.jpg
data/FingerDatasets/images/72.jpg
data/FingerDatasets/images/73.jpg
data/FingerDatasets/images/74.jpg
data/FingerDatasets/images/75.jpg
data/FingerDatasets/images/76.jpg
data/FingerDatasets/images/77.jpg
data/FingerDatasets/images/78.jpg
data/FingerDatasets/images/79.jpg
data/FingerDatasets/images/80.jpg
data/FingerDatasets/images/81.jpg

4.4 训练配置文件编写
在yolov5/data下新建一个FingerDetect.yaml文件。用于指引学习程序连接4.3 训练集测试集分配中的三个txt文件，连接数据集。内容如下：

train: data/FingerDatasets/train.txt
val:  data/FingerDatasets/val.txt
test: data/FingerDatasets/test.txt
 
nc: 1
 
names: ['FInger']

4.5 修改模型
本例我们使用v5x模型，打开模型文件yolov5/models/yolov5x.yaml 。对于复现，只需要修改开头几行，修改识别类型的数量，适当调整网络深度和宽度。

熟悉网络结构后可自行调整锚框、backbone、head。

主要的修改项为：

nc: 1  # number of classes
depth_multiple: 1.5  # model depth multiple
width_multiple: 1.33  # layer channel multiple

4.6 修改训练文件

打开yolov5/train.py , 在def parse_opt 函数中，根据实际情况修改训练参数。

主要的修改项为：
权重 - 选择预训练模型，连接到下载的v5x权重

parser.add_argument('--weights', type=str, default=ROOT / 'weights/yolov5x.pt', help='initial weights path')

模型 - 选择使用的模型，连接到yolov5x.yaml文件

parser.add_argument('--cfg', type=str, default='models/yolov5x.yaml', help='model.yaml path')

数据集 - 连接数据集的引导yaml文件

parser.add_argument('--data', type=str, default=ROOT / 'data/FingerDetect.yaml', help='dataset.yaml path')

EPoch数量 - 最大训练次数。默认最近100次训练无进展后自动停止。

parser.add_argument('--epochs', type=int, default=500)

Batch大小 - 训练块大小。过大可能会CUDA报错空间不足

parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs, -1 for autobatch')

4.7 开始训练

cd yolov5
python3 train.py

5. 训练结果

训练结果默认放置在yolov5/runs文件夹下最新的一个exp文件夹。

权重模型 - exp里生成的weights文件夹下
运行日志 - 名字最长的那个文件 events.out.tfevents.....
运行日志数据表 - results.csv
lose曲线 / precision曲线 - results.png
val集结果 - val_batchXX_labels.jpg

6.结果展示

可以用tensorboard可视化结果，例如以5.训练结果中的exp7为例：

cd yolov5/runs
tensorboard --logdir=exp7

在命令行输出的以下结果中右键打开localhost连接

即可在浏览器中看到tensorboard的显示结果