项目地址 https://github.com/ultralytics/yolov5

参考 https://zhuanlan.zhihu.com/p/711356735

目录

1  数据集准备

1.1  images

1.2  labels

1.3  yaml文件

2  环境配置

3  python环境配置

3.1  安装torch

3.2  安装opencv

3.3  安装 ultralytics

4  预训练模型准备

5  训练

6  预测

6.1  看大概的预测结果

6.2  预测命令

6.3  获取标签

---

 

1  数据集准备

我数据集单独放了一个文件夹，放在yolov5-master路径下

dataset中有两个文件夹 images 与 labels

1.1  images

images中有训练图像(train)与测试图像(val)，val与train的数量比2:8

我们打开train看一下

打开test看一下

1.2  labels

labels中有训练标注数据(train)与测试标注数据(val)，标注数据与图像数据一一对应（包括文件名）

我们打开train看一下

打开val看一下

标注文件txt是标准的yolo格式，我们随便打开一个txt看一下

1.3  yaml文件

在 yolov5-master/data 文件夹中创建custom.yaml文件，新建一个文本文档然后改尾缀就行了

内容如下

• path是dataset的绝对路径
• train是训练图片路径
• val是测试图片路径
• test可以不写
• nc是种类数，是几分类就写几
• names是标签，顺序最好与标注时一直（不一致也没事儿，后面可以再改，实际上训练出来的模型也不知道是哪个标签，它只认txt文件中的0、1、2这种序号）

2  环境配置

用的是亚马逊的EC2，型号是ml.t3.medium

一个小时需要0.2684元

这个型号是目前（2024.10.8）基本上可以说是性能最差的型号，应该是2核CPU的意思

如果要加速训练的话还是建议用这个

这个是将近34块钱一小时

虽然我的型号不支持GPU，但是执行nvcc -V还是能出东西

但是nvidia-smi出不了

查显卡的版本也查不了

3  python环境配置

我用的机器配置好了CUDA，所以我就直接配置python环境了，如果没配置好CUDA也可以尝试用CPU训练

我当前的python环境是3.10.8

3.1  安装torch

我安装pytorch的时候不知道我当前的云计算平台上用不了gpu，所以就按照CUDA装的，如果硬件没问题的话，安装方式是没问题的

可以在官网查看安装的命令 PyTorch

现在安装pytorch看起来简单的很，CUDA的版本它自己就给你搞定了

速度太慢了，否掉后用的清华源

清华源到后面一样会帮你解决cuda问题

最后安装成功是这样的

由于硬件问题，我这里如果使用 torch.cuda.is_available() 显示的是 false

• 会导致训练特别慢，大概要16个小时才能训练好一个模型

3.2  安装opencv

pip3 install opencv-python -i Simple Index

3.3  安装 ultralytics

pip3 install ultralytics -i Simple Index

4  预训练模型准备

预训练模型地址 Releases · ultralytics/yolov5 · GitHub

搜索pretrained，然后一直下一个到第五个（随着不断的更新可能就不是第5个了，不过我觉得使用别的应该也差不多，只要名字一样就行了）。点表格里的YOLOv5s是没用的

要到该条下面的Assets中找

下载之后的名称叫yolov5s.pt ，我把它的名字改成了pretrained_yolov5s.pt

5  训练

python train.py --weights pretrained_yolov5s.pt --data data/custom.yaml --workers 1 --batch-size 4

• 如果batch-size调整为8，我当前的机器会崩溃
• pretrained_yolov5s.pt本身也是一个模型，如果你把你训练好的pt文件放到这里就相当于是中继训练了

可能会出现少一个字体文件

在本地下载完 https://github.com/ultralytics/assets/releases/download/v0.0.0/Arial.ttf 然后上传上去，之后复制到 /home/ec2-user/.config/Ultralytics/Arial.ttf 就行了

之后再次执行 python train.py --weights pretrained_yolov5s.pt --data data/custom.yaml --workers 1 --batch-size 4

训练时会自动给你调整epoch，终端上显示的进度条走到100%的时候是一个epoch。进度条后面的时间比如 00：04<13:40 的意思是，该epoch已经训练了4秒，完成该epoch预计需要13分钟40秒

训练了19个epoch后我就否掉了，因为实在太慢了，否掉也是有模型出现的，在yolov5-master/runs中会出现文件夹train

打开train后会根据训练的次数出现exp，我训练了三次，所以出现了三个exp，最后一次训练是exp3

模型在 exp3 的 weights 中

其中best.pt是最好的，last.pt是最后的

6  预测

6.1  看大概的预测结果

我们可以通过 train_batch 的这三张图像大概看一下预测结果

我们发现在有些图像上还是可以的

6.2  预测命令

执行 python detect.py --weights ./runs/train/exp3/weights/best.pt --source ./dataset/images/val/road0.jpg 可以进行预测

执行后在 yolov5-master/runs 会出现 detect

里面有exp

• 预测几个就有几个exp

再里面有预测好的图像

打开看效果还行

6.3  获取标签

如果想获取标签的名称可以在detect中做一些修改，想这种写的比较好的项目，变量的命名都比较标准，在detect.py中搜索label

在这里给他打印出来

• 我估计下面的confience应该就是置信度

运行后发现可以在终端打印