数据集的训练和模型生成
准备数据集
因为我是使用GitHub上开源的数据集,用于红绿灯检测。刚开始下载了数据集后急着想要尝试一下在电脑上训练模型,但很快就被打脸了。
在windows上面跑pytorch代码,导致:
OSError: [WinError 1455] 页面文件太小,无法完成操作。 Error loading “D:\python\Anancoda\envs\torch38\lib\site-packages\torch\lib\caffe2_detectron_ops_gpu.dll” or one of its dependencies.
,在CSDN上查看解决问题的方法,
【pytorch问题解决】OSError: [WinError 1455] 页面文件太小,无法完成操作。
按照这篇文章的操作就解决了。
因为我是直接使用开源的数据集,里面的标注文件是.xml格式,而yolov5训练时的标注文件必须得是.txt文件,所有我还需要写脚本更改文件的格式。
#本文件用于将数据集中的xml文件转换为txt文件,用于yolo训练
#代码只需要修改三处位置即可
import os
import glob
import xml.etree.ElementTree as ET
#xml文件存放位置
xml_file=r'E:\K510\yolov5-6.0\xml1'
l=['red', 'off','yellow','green']
def convert(box,dw,dh):
x=(box[0]+box[2])/2.0
y=(box[1]+box[3])/2.0
w=box[2]-box[0]
h=box[3]-box[1]
x=x/dw
y=y/dh
w=w/dw
h=h/dh
return x,y,w,h
def f(name_id):
xml_o=open(r'E:\K510\yolov5-6.0\xml1\%s.xml'%name_id)#xml文件存放的位置
txt_o=open(r'E:\K510\yolov5-6.0\database\train\labels\%s.txt'%name_id,'w')#txt文件存放的位置
pares=ET.parse(xml_o)
root=pares.getroot()
objects=root.findall('object')
size=root.find('size')
dw=int(size.find('width').text)
dh=int(size.find('height').text)
for obj in objects :
c=l.index(obj.find('name').text)
bnd=obj.find('bndbox')
b=(float(bnd.find('xmin').text),float(bnd.find('ymin').text),
float(bnd.find('xmax').text),float(bnd.find('ymax').text))
x,y,w,h=convert(b,dw,dh)
write_t="{} {:.5f} {:.5f} {:.5f} {:.5f}\n".format(c,x,y,w,h)
txt_o.write(write_t)
xml_o.close()
txt_o.close()
name=glob.glob(os.path.join(xml_file,"*.xml"))
for i in name :
name_id=os.path.basename(i)[:-4]
f(name_id)
这个python代码如果要直接使用的话只需要更改一下各文件的地址和类别的名称。
注意:在转文件时需要注意类别的顺序,以免到时候训练出来的模型标签混乱。
当时还修改了一下/yolov5-6.0/utils/plots.py文件,在函数定义的下端定义了一个新的数组,用于存放类别,又在最后的cv2.putText()函数中修改了label的参数。
当时遇到这个问题时因为实在不知道如何解决,只能够尝试着先训练10次查看导出的模型识别到的结果正不正确,但最后瞎猫碰到死耗子被我给解决了。
这时标注数据也算是正常了。
但是在训练前还需要将数据集分成训练集和验证集。在模型训练中,需要有训练集和验证集。可以简单理解为网络使用训练集去训练,训练出来的网络使用验证集验证。在总数据集中训练集通常应占80%,验证集应占20%。
最后在yolov5-6.0\data目录下创建一个新的data.yaml文件,打开这个文件去指定需要训练的训练数据集和验证数据集,以及我们的类别数和类别名称,(标签混乱的问题也可能是这里的代码顺序写错了,但是我没有改这个地方的顺序,之后如果再次遇到这个问题可以考虑一下)。
在yolov5-6.0\models目录下我们也有创建一个新的yaml文件,用来修改类别数
我们只需要在这里修改参数nc的值。
我们还要在yolov5-6.0目录下打开train.py文件
把cfg参数的值改为上一步在yolov5-6.0\models目录下我们创建yaml文件的名称,data参数的值修改为data目录下创建的文件的名称类似于这样
parser.add_argument('--cfg', type=str, default=ROOT / 'models/my_model.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/data.yaml', help='dataset.yaml path')
如果需要修改训练的轮次,可以修改在train.py文件下的epochs的值
最后的最后,打开Anaconda Prompt运行台,输入conda env list,查看电脑上的环境是否搭建好,需要的是py37_yolov5环境。
输入conda activate py37_yolov5 激活该环境,进入yolov5-6.0源码的目录,进入之后输入python train.py即可开始训练。
若要自行准备数据集,需要下载标注工具
LabelImg标注工具下载地址
数据集目录需要我们自行创建images文件夹和labels文件夹分别存放图像数据集和标注信息。自主标注的数据集标注完成后只需要修改一下yolo源码中数据集的地址和类别这些参数,即可开始训练。