【目标检测】yolov5训练自己的数据【数据准备、训练、样本预测】
前言：
- 安装Anaconda：
- 创建虚拟环境：
- 安装pytorch：
- 下载源码和安装依赖库：
- 数据标注：
- 数据预处理：
- 下载预训练模型： 7.训练 8.测试
一、前言 
下载完成后打开一路Yes即可，只需要注意这里要将conda添加到PATH：
二. 创建虚拟环境：
这里我们需要为yolov5单独创建一个环境，输入：
conda create -n torch107 python=3.7
安装完成后，输入：
activate torch107
激活环境： 注：删除某个虚拟环境命令：
conda remove -n your_env_name(虚拟环境名称) --all
三. 安装pytorch：
yolov5最新版本需要pytorch1.6版本以上，因此我们安装pytorch1.7版本。由于我事先安装好了CUDA10.1，因此在环境中输入，即可安装：：
pip install torch==1.7.0+cu101 torchvision==0.8.1+cu101 torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
然后查看CUDA是否可用：
四. 下载源码和安装依赖库： 源码地址：https://github.com/ultralytics/yolov5
下载后解压，在目录内打开cmd并激活环境：
安装依赖库：
pip install -r requirements.txt
5. 数据标注： 数据标注我们要用labelimg，使用pip即可安装：
pip install labelimg
六、数据准备
将图像和数据统一放置到源码目录的VOCData文件夹下。 其中，jpg文件放置在VOCData/images下，xml放置在VOCData/Annotations下：
数据预处理： 创建 split.py 文件，内容如下：
import os
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='VOCData/Annotations', type=str, help='input xml label path')
parser.add_argument('--txt_path', default='VOCData/labels', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
for i in list_index:
name = total_xml[i][:-4] + '\n'
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
运行结束后，可以看到VOCData/labels下生成了几个txt文件：
然后新建 txt2yolo_label.py 文件用于将数据集转换到yolo数据集格式：
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
from tqdm import tqdm
import os
from os import getcwd
sets = ['train', 'val', 'test']
classes = ['face', 'normal', 'phone', 'write',
'smoke', 'eat', 'computer', 'sleep']
def convert(size, box):
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return x, y, w, h
def convert_annotation(image_id):
# try:
in_file = open('VOCData/Annotations/%s.xml' % (image_id), encoding='utf-8')
out_file = open('VOCData/labels/%s.txt' % (image_id), 'w', encoding='utf-8')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
b1, b2, b3, b4 = b
# 标注越界修正
if b2 > w:
b2 = w
if b4 > h:
b4 = h
b = (b1, b2, b3, b4)
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " +
" ".join([str(a) for a in bb]) + '\n')
# except Exception as e:
# print(e, image_id)
wd = getcwd()
for image_set in sets:
if not os.path.exists('VOCData/labels/'):
os.makedirs('VOCData/labels/')
image_ids = open('VOCData/labels/%s.txt' %
(image_set)).read().strip().split()
list_file = open('VOCData/%s.txt' % (image_set), 'w')
for image_id in tqdm(image_ids):
list_file.write('VOCData/images/%s.jpg\n' % (image_id))
convert_annotation(image_id)
list_file.close()
转换后可以看到VOCData/labels下生成了每个图的txt文件：
在data文件夹下创建myvoc.yaml文件：
内容如下：
train: VOCData/train.txt
val: VOCData/val.txt
# number of classes
nc: 8
# class names
names: ["face", "normal", "phone", "write", "smoke", "eat", "computer", "sleep"]
七. 下载预训练模型： 我训练yolov5m这个模型，因此将它的预训练模型下载到weights文件夹下，下载地址如下： https://github.com/ultralytics/yolov5/releases
修改models/yolov5m.yaml下的类别数：
八、训练train.py
train.py中的一些参数修改 最后，在根目录中对train.py中的一些参数进行修改，主要参数解释如下。我们平时训练的话，主要用到的只有这几个参数而已：–weights，–cfg，–data，–epochs，–batch-size，–img-size，–project。
parser = argparse.ArgumentParser()
# 加载的权重文件
parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
# 模型配置文件，网络结构，使用修改好的yolov5m.yaml文件
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
# 数据集配置文件，数据集路径，类名等，使用数据集方面的cat.yaml文件
parser.add_argument('--data', type=str, default='data/cat.yaml', help='data.yaml path')
# 超参数文件
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
# 训练总轮次，1个epoch等于使用训练集中的全部样本训练一次，值越大模型越精确，训练时间也越长。
parser.add_argument('--epochs', type=int, default=300)
# 批次大小，一次训练所选取的样本数，显卡垃圾的话，就调小点
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
# 输入图片分辨率大小，nargs='+'表示参数可设置一个或多个
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
# 是否采用矩形训练，默认False，开启后可显著的减少推理时间
parser.add_argument('--rect', action='store_true', help='rectangular training')
# 接着打断训练上次的结果接着训练
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
# 不保存模型，默认False
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
# 不进行test，默认False
parser.add_argument('--notest', action='store_true', help='only test final epoch')
# 不自动调整anchor，默认False
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
# 是否进行超参数进化，默认False
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
# 谷歌云盘bucket，一般不会用到
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
# 是否提前缓存图片到内存，以加快训练速度，默认False
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
# 选用加权图像进行训练
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
# 训练的设备，cpu；0(表示一个gpu设备cuda:0)；0,1,2,3(多个gpu设备)。值为空时，训练时默认使用计算机自带的显卡或CPU
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# 是否进行多尺度训练，默认False
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
# 数据集是否只有一个类别，默认False
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
# 是否使用adam优化器
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
# 是否使用跨卡同步BN,在DDP模式使用
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
# gpu编号
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
# W&B记录的图像数，最大为100
parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100')
# 记录最终训练的模型，即last.pt
parser.add_argument('--log-artifacts', action='store_true', help='log artifacts, i.e. final trained model')
# dataloader的最大worker数量
parser.add_argument('--workers', type=int, default=4, help='maximum number of dataloader workers')
# 训练结果所存放的路径，默认为runs/train
parser.add_argument('--project', default='runs/train', help='save to project/name')
# 训练结果所在文件夹的名称，默认为exp
parser.add_argument('--name', default='exp', help='save to project/name')
# 若现有的project/name存在，则不进行递增
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args()
https://blog.csdn.net/oJiWuXuan/article/details/107558286
注意： workers设置多少才合适，这个问题是很难有一个推荐的值。有以下几个建议：
num_workers=0表示只有主进程去加载batch数据，这个可能会是一个瓶颈。 num_workers = 1表示只有一个worker进程用来加载batch数据，而主进程是不参与数据加载的。这样速度也会很慢。 num_workers>0 表示只有指定数量的worker进程去加载数据，主进程不参与。增加num_works也同时会增加cpu内存的消耗。所以num_workers的值依赖于 batch size和机器性能。 一般开始是将num_workers设置为等于计算机上的CPU数量 最好的办法是缓慢增加num_workers，直到训练速度不再提高，就停止增加num_workers的值
运行命令：
python train.py --img 640 --batch 4 --epoch 300 --data ./data/myvoc.yaml --cfg ./models/yolov5m.yaml --weights weights/yolov5m.pt --workers 0
可能会报错如下，原因是gpu内存不足：
运行命令查看内存，输入 nvidia-smi，会显示GPU的使用情况，以及占用GPU的应用程序
解决方法： 1、batch size 设置小一点 2、设置其他gpu，如在代码中添加：
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
3、Linux下释放GPU显存
fuser -v /dev/nvidia*
查看当前系统中GPU占用的线程
使用kill -9 线程号释放显存
成功运行!
在训练中，也可以随时查看每一轮次训练的结果，可利用tensorboard可视化训练过程，训练开始时会在runs/train/exp文件夹中产生一个“events.out.tfevents.1608924773.JWX.5276.0”文件，利用tensorboard打开即可查看训练日志。首先我们通过cmd进去该YOLOv5所在的项目文件夹，然后激活所用的虚拟环境，输入如下命令行：
tensorboard --logdir runs/train/exp
九、测试detect.py 训练完成后，用自己的权重测试图片
有了训练好的权重后，就可以就行目标检测测试了。直接在根目录的detect.py中进行调试，主要参数解释如下。我们平时用的话，主要用到的只有这几个参数而已：–weights，–source，–conf-thres，–project。
parser = argparse.ArgumentParser()
# 选用训练的权重，可用根目录下的yolov5s.pt，也可用runs/train/exp/weights/best.pt
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
# 检测数据，可以是图片/视频路径，也可以是'0'(电脑自带摄像头),也可以是rtsp等视频流
parser.add_argument('--source', type=str, default='inference/videos/猫猫识别.mp4', help='source') # file/folder, 0 for webcam
# 网络输入图片大小
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
# 置信度阈值，检测到的对象属于特定类（狗，猫，香蕉，汽车等）的概率
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
# 做nms的iou阈值
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
# 检测的设备，cpu；0(表示一个gpu设备cuda:0)；0,1,2,3(多个gpu设备)。值为空时，训练时默认使用计算机自带的显卡或CPU
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# 是否展示检测之后的图片/视频，默认False
parser.add_argument('--view-img', action='store_true', help='display results')
# 是否将检测的框坐标以txt文件形式保存，默认False
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
# 是否将检测的labels以txt文件形式保存，默认False
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
# 设置只保留某一部分类别，如0或者0 2 3
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
# 进行nms是否也去除不同类别之间的框，默认False
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
# 推理的时候进行多尺度，翻转等操作(TTA)推理
parser.add_argument('--augment', action='store_true', help='augmented inference')
# 如果为True，则对所有模型进行strip_optimizer操作，去除pt文件中的优化器等信息，默认为False
parser.add_argument('--update', action='store_true', help='update all models')
# 检测结果所存放的路径，默认为runs/detect
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
# 检测结果所在文件夹的名称，默认为exp
parser.add_argument('--name', default='exp', help='save results to project/name')
# 若现有的project/name存在，则不进行递增
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args()
(1) 单张图片测试
python detect.py --source demo/test.jpg --weights runs/train/exp3/weights/best.pt
（2）批量图片测试
python detect.py --source ./demo --weights runs/train/exp3/weights/best.pt
（3）视频测试
python detect.py --source demo/demo.mp4 --weights weights/yolov5m.pt
十、经过训练的 YOLOv5 🚀 模型从 PyTorch 导出为 ONNX 和 TorchScript 格式
在你开始前 克隆这个 repo 并安装requirements.txt依赖项，包括Python>=3.8和PyTorch==1.7。
git clone https://github.com/ultralytics/yolov5 # clone repo
cd yolov5
pip install -r requirements.txt # base requirements
pip install coremltools>=4.1 onnx>=1.9.0 scikit-learn==0.19.2 # export requirements
导出经过训练的 YOLOv5 模型 此命令将预训练的 YOLOv5s 模型导出为 ONNX、TorchScript 和 CoreML 格式。yolov5s.pt是最轻、最快的型号。其他选项是yolov5m.pt, yolov5l.ptand yolov5x.pt, 或者您拥有训练自定义数据集的检查点runs/exp0/weights/best.pt。有关所有可用模型的详细信息
python export.py --weights runs/train/exp3/weights/best.pt --img 640 --batch 1
【小白CV】手把手教你用YOLOv5训练自己的数据集（从Windows环境配置到模型部署） RuntimeError: CUDA out of memory（已解决） tensorflow中使用指定的GPU及GPU显存 CUDA_VISIBLE_DEVICES YOLOv5说明文档 YOLOV5测试及训练自己的数据集
投递vivo等公司10个岗位
