前言
之前文章已经讲过yolov5模型的训练,这一篇将说一下分类模型训练流程。
https://blog.csdn.net/qq_45066628/article/details/129470290
新版本简介
YOLOv5官方发布了v6.2版本,v6.2版本支持分类模型训练、验证、预测和导出;v6.2版本的推出使得训练分类器模型变得超级简单!
v6.2版本项目结构并无太大改变,主要是增加了classify文件夹以及predict.py train.py val.py 这三个文件;那么这三个文件也分别对应着分类模型的推理、训练和验证。
训练前准备工作🌟
模型下载
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -qr requirements.txt
数据准备
这里数据集采用kaggle猫狗大战数据集
数据集下载地址:https://www.kaggle.com/c/dogs-vs-cats-redux-kernels-edition/data
划分数据集
分类模型无需标注数据,只需要将训练的图片按类别文件夹划分即可
格式如下🍀:
示例:
划分脚本:
import os
from shutil import copy
import random
def mkfile(file):
if not os.path.exists(file):
os.makedirs(file)
# 获取data文件夹下所有文件夹名(即需要分类的类名)
file_path = '../test_set/'
flower_class = [cla for cla in os.listdir(file_path)]
# 创建 训练集train 文件夹,并由类名在其目录下创建5个子目录
mkfile('data/train')
for cla in flower_class:
mkfile('data/train/' + cla)
# 创建 验证集val 文件夹,并由类名在其目录下创建子目录
mkfile('data/val')
for cla in flower_class:
mkfile('data/val/' + cla)
# 划分比例,训练集 : 验证集 = 8 : 2
split_rate = 0.2
# 遍历所有类别的全部图像并按比例分成训练集和验证集
for cla in flower_class:
cla_path = file_path + '/' + cla + '/' # 某一类别的子目录
images = os.listdir(cla_path) # iamges 列表存储了该目录下所有图像的名称
num = len(images)
eval_index = random.sample(images, k=int(num * split_rate)) # 从images列表中随机抽取 k 个图像名称
for index, image in enumerate(images):
# eval_index 中保存验证集val的图像名称
if image in eval_index:
image_path = cla_path + image
new_path = 'data/val/' + cla
copy(image_path, new_path) # 将选中的图像复制到新路径
# 其余的图像保存在训练集train中
else:
image_path = cla_path + image
new_path = 'data/train/' + cla
copy(image_path, new_path)
print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="") # processing bar
print()
print("processing done!")
下载权重文件
下载地址:https://github.com/ultralytics/yolov5
不同文件,使用产生的效果也不一样,这里使用YOLOv5s-cls
修改配置文件
到classify文件夹下,修改train.py
第一行设置自己下载的权重文件路径,第二行设置数据集路径,其余参数可根据自行需要修改。
开始训练🚀
-
方法一:
运行classify下的train.py文件,这里由于cuda问题,暂时选用cpu进行训练演示。
运行后这里缺失模块产生了一个报错,若没有跳过,解决方法:pip install --upgrade protobuf
然后重新运行train.py文件
-
方法二:
python classify/train.py --model yolov5s-cls.pt --data datasets/data–epochs 100 --batch-size 32 --imgsz 224
训练完成后,会在runs文件夹下看到模型
验证和推理
验证
-
方法一:
python classify/val.py --weights runs/train-cls/exp4/weights/best.pt --data datasets/data
-
方法二:
到classify文件夹下,修改val.py
第一行使用自己训练的模型文件,第二行修改数据集路径,其余参数可根据自行需要修改。
验证结果:
输出类别和正确率信息
推理
方法与上述类似,该文件或直接使用命令行
# 测试im1.jpg
python classify/predict.py --weights runs/train-cls/exp4/weights/best.pt --source im1.jpg
# 测试im2.jpg
python classify/predict.py --weights runs/train-cls/exp4/weights/best.pt --source im2.jpg
推理结果:
可以到runs文件夹下的predict-cls文件夹下查看结果
导出
使用ONNX
执行命令导出onnx:
python export.py --weights runs/train-cls/exp4/weights/best.pt --include onnx
输出:
Detect: python classify/predict.py --weights runs/train-cls/exp4/weights/best.onnx
Validate: python classify/val.py --weights runs/train-cls/exp4/weights/best.onnx
PyTorch Hub: model = torch.hub.load(‘ultralytics/yolov5’, ‘custom’, ‘runs/train-cls/exp4/weights/best.onnx’) # WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference
Visualize: https://netron.app
使用TensorRT
执行命令导出engine:
python export.py --weights runs/train-cls/exp4/weights/best.pt --include engine --device 0
输出:
Detect: python classify/predict.py --weights runs/train-cls/exp4/weights/best.engine
Validate: python classify/val.py --weights runs/train-cls/exp4/weights/best.engine
PyTorch Hub: model = torch.hub.load(‘ultralytics/yolov5’, ‘custom’, ‘runs/train-cls/exp/weights/best.engine’) # WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference
Visualize: https://netron.app