yolov5s 预训练模型_使用yolo-v5训练测试自己的数据

使用环境：python3.6.8、CUDA10.0、Centos7.5

一、前言

1.1下载yolov5代码

首先通过git命令从GitHub克隆yolov5最新代码(当然也可以选择想要的版本号)，该文采用最新的代码，命令如下：

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

其中 --recursive 参数属于个人习惯添加，主要是防止项目中参杂第三方的工程项目；克隆完成以后的文件夹如下(以下显示的是克隆到本地的全部内容，可能有些.git、.dockerignore等等文件因为是隐藏文件没有显示出来，不必在意)：

1.2安装yolov5运行时需要的包

在克隆的代码中可以找到requirements.txt文件，为了避免踩坑，尽量不要直接安装里面的包，主要是针对pytorch版本；

如果你的电脑没有GPU，可以直接运行以下命令：

pip3 install -r requirements.txt

如果你的电脑有GPU，请不要直接运行以上命令，先更新pytorch为1.6.0，具体命令如下：

pip3 install torch==1.6.0+cu92 torchvision==0.7.0+cu92 -f https://download.pytorch.org/whl/torch_stable.html

然后把requirements.txt中有关torch的内容删除，删除之后的requirements.txt文件为：

而后在运行：

pip3 install -r requirements.txt

1.3下载预训练模型和测试

下载预训练模型时尽量从yolov5给定的链接处下载；

之所以不建议从别的地方下载预训练模型，原因是由于时间关系有可能下载到的模型和yolov5的代码不匹配，导致各种乱七八糟的错误，因此最好从上述链接上下载；(本文下载的是yolov5x.pt，如果因为网络原因导致模型下载失败，也没有关系，直接跳过测试就好，不影响训练自己的数据模型)

模型下载完成以后放在和detect.py同级的目录下，大致如下：

将detect.py中的代码大致150~164行代码做以下修改：

parser = argparse.ArgumentParser()

parser.add_argument('--weights', nargs='+', type=str, default='yolov5x.pt', help='model.pt path(s)') # 本文下载的是yolov5x.pt

parser.add_argument('--source', type=str, default='inference/images', help='source')

parser.add_argument('--output', type=str, default='inference/output', help='output folder')

parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')

parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')

parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')

parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') # 没有GPU的话请改为 default='cpu'

parser.add_argument('--view-img', action='store_true', help='display results')

parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')

parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')

parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')

parser.add_argument('--augment', action='store_true', help='augmented inference')

parser.add_argument('--update', action='store_true', help='update all models')

opt = parser.parse_args()

之后运行detect.py，结果如下：

效果一级棒！

二、制作自己的训练数据集

2.1使用LabelImg标注自己的数据

yolov5训练集使用的是yolo的格式的标注方式，具体格式如下图所示：

每行数据代表一个物体，每一行的含义为：class_id center_x center_y width height，一个txt文件对应一个图片，例如上图000000000009.txt对应的既是000000000009.jpg，标注时只需将LabelImg中的 PascalVOC点击转换为YOLO即可；

2.2从VOC标注数据转换为YOLO标注格式

如果已经有了标注出来的VOC标注文件(xml格式文件)，只需要使用脚本转换为YOLO格式即可，这种方式能最快把数据集准备好，代码如下：

"""

Creation Time : 2020.08.14

Create author : Tian FuKang

"""

import os

import glob

import pandas as pd

import random

import cv2

import xml.etree.ElementTree as ET

def resize_image(img, max_l):

w = img.shape[1]

h = img.shape[0]

if w > h:

scale = float(max_l) / float(w)

w = max_l

h = int(h * scale)

else:

scale = float(max_l) / float(h)

h = max_l

w = int(w * scale)

img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)

return img, scale

def xml_to_csv(path_1):

xml_file_list = []

# 在这里修改你需要检测的类别名字

class_text = ['person', 'bicycle', 'car', 'motorcycle', 'airplane']

for xml_file in glob.glob(path_1 + '/*.xml'):

xml_file_list.append(xml_file)

for i in range(0, len(xml_file_list)):

tree = ET.parse(xml_file_list[i])

root = tree.getroot()

filename = root.find('filename').text

name = str(i + 1).zfill(12)

# 在这里修改保存数据和标签的文件路径

img_name = './data/images/train/' + name + '.jpg'

txt_name = './data/labels/train/' + name + '.txt'

file_txt = open(txt_name, mode='a', encoding='utf-8')

image = cv2.imread(path_1.replace('image_xml', 'image/') + filename)

image, scale = resize_image(image, 640)

cv2.imwrite(img_name, image)

for member in root.findall('object'):

class_id = class_text.index(member[0].text)

xmin = int(member[4][0].text)

ymin = int(member[4][1].text)

xmax = int(member[4][2].text)

ymax = int(member[4][3].text)

# round(x, 6) 这里我设置了6位有效数字，可根据实际情况更改

center_x = round(((xmin + xmax) / 2.0) * scale / float(image.shape[1]), 6)

center_y = round(((ymin + ymax) / 2.0) * scale / float(image.shape[0]), 6)

box_w = round(float(xmax - xmin) * scale / float(image.shape[1]), 6)

box_h = round(float(ymax - ymin) * scale / float(image.shape[0]), 6)

file_txt.write(str(class_id))

file_txt.write(' ')

file_txt.write(str(center_x))

file_txt.write(' ')

file_txt.write(str(center_y))

file_txt.write(' ')

file_txt.write(str(box_w))

file_txt.write(' ')

file_txt.write(str(box_h))

file_txt.write('\n')

file_txt.close()

def main():

# 在这里修改xml文件的路径，我的xml文件在image_xml文件夹中，因此是以下设置方式

path_2 = './image_xml'

xml_to_csv(path_2)

print('Successfully converted xml to csv.')

main()

运行该脚本时需要注意以下几点：1.保证修改了第36行，更改为自己的类别名称；2.保证第50、51行的保存数据的文件夹存在；3.保证你的数据图片和标注文件分别在image文件夹和image_xml文件夹里面，并且image文件夹和image_xml文件夹在同一目录下，因为在第55、90行用到了这两个文件夹；之后运行该脚本即可；

三、使用yolov5训练自己的数据集

3.1规范自己的数据集

进入到之前目录1.1中克隆的代码的文件夹中，找到data文件夹，在data文件夹下新建如下目录文件夹；

其中目录：

./data/custom/images/train/ 存放训练集图片

./data/custom/images/val/ 存放验证集图片

./data/custom/labels/train/ 存放训练集标签

./data/custom/labels/val/ 存放验证集标签

将自己上个步骤生成的数据集按照训练集验证集放入以上目录中(训练集验证集数据比例可自己控制)；注意：训练集中的image和label的文件名要一一对应(例如./data/custom/images/train/0001.jpg和./data/custom/labels/train/0001.txt)，同理验证集也是如此；

3.2编写yaml数据配置文件

还是在yolov5项目的data文件夹下，新建文件custom.yaml，如下图所示：

编辑内容如下：

# 设置训练集和验证集的目录

train: ./data/code/images/train/

val: ./data/code/images/val/

# 类别数目

nc: 5

# 类别名

names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane']

保存即可；

3.3修改yolov5*.yaml配置文件

本文使用的yolov5s，因此修改的是yolov5s.yaml，修改方式很简单，只需要修改类别数目即可，修改后的内容如下：

# parameters

nc: 5 # 主要修改此处为自己的类别数量

depth_multiple: 0.33 # model depth multiple

width_multiple: 0.50 # layer channel multiple

# anchors

anchors:

- [10,13, 16,30, 33,23] # P3/8

- [30,61, 62,45, 59,119] # P4/16

- [116,90, 156,198, 373,326] # P5/32

# YOLOv5 backbone

backbone:

# [from, number, module, args]

[[-1, 1, Focus, [64, 3]], # 0-P1/2

[-1, 1, Conv, [128, 3, 2]], # 1-P2/4

[-1, 3, BottleneckCSP, [128]],

[-1, 1, Conv, [256, 3, 2]], # 3-P3/8

[-1, 9, BottleneckCSP, [256]],

[-1, 1, Conv, [512, 3, 2]], # 5-P4/16

[-1, 9, BottleneckCSP, [512]],

[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32

[-1, 1, SPP, [1024, [5, 9, 13]]],

[-1, 3, BottleneckCSP, [1024, False]], # 9

]

# YOLOv5 head

head:

[[-1, 1, Conv, [512, 1, 1]],

[-1, 1, nn.Upsample, [None, 2, 'nearest']],

[[-1, 6], 1, Concat, [1]], # cat backbone P4

[-1, 3, BottleneckCSP, [512, False]], # 13

[-1, 1, Conv, [256, 1, 1]],

[-1, 1, nn.Upsample, [None, 2, 'nearest']],

[[-1, 4], 1, Concat, [1]], # cat backbone P3

[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)

[-1, 1, Conv, [256, 3, 2]],

[[-1, 14], 1, Concat, [1]], # cat head P4

[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)

[-1, 1, Conv, [512, 3, 2]],

[[-1, 10], 1, Concat, [1]], # cat head P5

[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)

[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)

]

3.4修改运行train.py文件

完成上述步骤以后，还需要修改train.py文件中的一些代码，主要集中在第383-409行，修改后如下所示：

parser = argparse.ArgumentParser()

# --weights 因为是训练自己的数据集，因此没有必要加载预训练模型，这里改成无，当然也可以加载

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

# --cfg 这里改成默认是./models/yolov5x.yaml,因为使用了yolov5s

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

# --data 这里改为默认是./data/custom.yaml

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

parser.add_argument('--hyp', type=str, default='', help='hyperparameters path, i.e. data/hyp.scratch.yaml')

# --epochs 迭代轮次可以修改，本文默认为300

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

# --batch-size 硬件不好的话，把16改为8或者4甚至更低(yolov5支持更小，问题不大)

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

# --img-size 图片大小可以根据自己的数据设置，注意测试时同步大小就好，这里不做修改

parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes')

parser.add_argument('--rect', action='store_true', help='rectangular training')

parser.add_argument('--resume', nargs='?', const='get_last', default=False,

help='resume from given path/last.pt, or most recent run if blank')

parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')

parser.add_argument('--notest', action='store_true', help='only test final epoch')

parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')

parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')

parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')

parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')

parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')

# --device 我的电脑有GPU卡，因此设置为0，没有设置为cpu

parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')

parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')

parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')

parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')

parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')

parser.add_argument('--logdir', type=str, default='runs/', help='logging directory')

parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')

opt = parser.parse_args()

修改完以后运行train.py，运行过程如下：

训练完成以后会在项目中生成以下内容：

至此训练完成；

3.5测试生成的模型

快速测试只需要修改detect.py文件中的第164-178行即可，修改跟训练时类似，如下所示：

parser = argparse.ArgumentParser()

parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')

parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam

parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder