当前位置：首页 > news >正文

如何使用yolov5-master进行训练

news 来源：原创 2024/9/20 12:39:17

主要记录配置完成的yolov5-master如何进行自定义数据集训练。

本人使用的设备是jetson-orin AGX，系统是英伟达官方魔改过的ubuntu18.04.。

一、修改配置

首先观察train.py的__main__部分，我们发现训练程序接受很多个参数。

其中我们必须要关注的参数包括：

1、 --weight 初始权重设置，需要输入初始权重的路径，但也可以设置为 '' , 使得模型从零开始训练

2、 --cfg 模型配置设置，需要输入模型初始配置文件的地址

3、 --data 数据配置设置，需要输入数据配置文件的地址

4、 --device 使用何种设备训练？设置为0，便可以使用gpu训练

1、配置--cfg

如下所示，模型配置文件一般储存在models文件夹里。

打开后发现里面已经有yolov5s.yaml，yolov5m.yaml，yolov5l.yaml和yolov5x.yaml四个文件。模型配置文件储存的是神经网络的架构信息，譬如，第几层是什么样的卷积层，第几层是什么样的全连接层。而这四个文件，便对应着四种复杂度从低到高的配置，你可以选择一个复制出来，像我一样命名为myyolo.yaml（或者是其他名称）。

打开这个配置文件，发现中间确实存着各种各样的架构信息。对于我这样的小白，这些都不重要，重要的是第一行nc这个参数，需要进行修改。

nc是什么呢？就是这个模型要识别多少种物体。比如你要识别人、电视、椅子这三种对象，就把它设置为3就行了

2、配置--data

打开主文件夹下的data文件夹，发现里面已经有很多配置文件了。打开coco128.yaml，我们查看一下其内容。

我们可以复制一个coco128.yaml然后按照自己的需求修改，也可以直接改（train.py默认使用这个配置文件）

注意看，train和val后分别保存着训练数据集和验证数据集的地址（这个地址是相对train.py这个文件的）；nc是识别物体种类数，names就是各个种类对应名称。

修改这些配置后就可以了，download不用管。

二、准备数据集

网上搜一搜coco128.zip这个文件，也可以点这个：

coco128数据集下载连接-CSDN博客

把文件解压在在coco128.yaml中配置的位置。打开后发现如下图所示，图片和数据是用12位数字命名的，而且数据和标记的主名称必须一样！

替换这些图片和数据，就可以使用自训练数据。

打开任意一个标记文件。发现它的格式如下所示：

分别标记着物体对应标号（可以看看coco128.yaml，'person'的标号应该是0，也就是说，标号是在data配置文件中设置的）

而后面四个浮点数，应该分别是：

x_center 横向坐标

y_center 纵向坐标

width 识别框宽度

height 识别框高度

当然，对这些数据进行手动标记太耗费时间，使用LabelImg这样的工具可以解决这个问题。这里不过多描述，如果只是像我一样的小白熟悉框架，可以使用coco128训练集！

三、开始训练

开始运行train.py

python train.py --data ./data/coco128.yaml --cfg ./models/myyolo.yaml --weights '' --device 0

输入我们刚才提及的一些系列参数配置，训练就开始了！如果一切顺利，那么皆大欢喜，但往往会遇到不少bug，这边建议联合gpt和百度搜索解决。

以下是几个我遇到的问题和解决方法：

1、np.int错误

由于我们装的numpy是新版本，弃用了np.int这个类型，改用np.int64；而yolov5考虑的是老版本的numpy，所以会有不兼容的地方。

所以建议安装1.20版本以下的numpy，或者就只能手动把报错的np.int改为np.int64

2、pillow错误

YOLOV5训练过程中出现报错：module ‘PIL.Image‘ has no attribute ‘ANTIALIAS‘_image.antialias-CSDN博客

pillow版本太高也有问题，这里有对应解决办法。

3、tensor不能直接转nparray

TypeError: can‘t convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to_return self.numpy() typeerror: can't convert cuda:-CSDN博客

4、cache错误

yolov7 magic_number = pickle_module.load(f, **pickle_load_args) _pickle.UnpicklingError: STACK_G_magic number = pickle module.load(f **pickle load -CSDN博客

下载的coco128.zip原有的cache文件必须全部删去