Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换

前言

之前文章已经讲过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 这三个文件;那么这三个文件也分别对应着分类模型的推理、训练和验证。
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第1张图片

训练前准备工作

模型下载

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -qr requirements.txt

数据准备

这里数据集采用kaggle猫狗大战数据集
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第2张图片
数据集下载地址:https://www.kaggle.com/c/dogs-vs-cats-redux-kernels-edition/data
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第3张图片

划分数据集

分类模型无需标注数据,只需要将训练的图片按类别文件夹划分即可

格式如下:

Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第4张图片
示例:
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第5张图片
划分脚本:

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/'
save_path = './data/'
flower_class = [cla for cla in os.listdir(file_path)]




# 划分比例,训练集 : 测试集 :验证= 7:1:2
split_train = 0,7
split_val = 0.1
split_test = 0.2

# 遍历所有类别的全部图像并按比例分成训练集和验证集
for cla in flower_class:
    cla_path = file_path + '/' + cla + '/'  # 某一类别的子目录
    images = os.listdir(cla_path)  # iamges 列表存储了该目录下所有图像的名称
    num = len(images)
    eval_train = random.sample(images, k=int(num * split_train))  # 从images列表中随机抽取 k 个图像名称
    eval_val = random.sample(images, k=int(num * split_val))  

    for index, image in enumerate(images):
    
        if image in eval_train:
            image_path = cla_path + image
            new_path = save_path+'/tarin/' + cla
            mkfile(new_path)
            copy(image_path, new_path)  # 将选中的图像复制到新路径

        elif image in eval_val:
            image_path = cla_path + image
            new_path = save_path+'/val/' + cla
            mkfile(new_path)
            copy(image_path, new_path)  # 将选中的图像复制到新路径

        else:
            image_path = cla_path + image
            new_path = save_path+'/test/' + cla
            mkfile(new_path)
            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
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第6张图片
不同文件,使用产生的效果也不一样,这里使用YOLOv5s-cls

修改配置文件

到classify文件夹下,修改train.py
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第7张图片
第一行设置自己下载的权重文件路径,第二行设置数据集路径,其余参数可根据自行需要修改。

开始训练

  1. 方法一:
    运行classify下的train.py文件,这里由于cuda问题,暂时选用cpu进行训练演示。
    Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第8张图片
    运行后这里缺失模块产生了一个报错,若没有跳过,解决方法:

    pip install --upgrade protobuf

    然后重新运行train.py文件
    Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第9张图片

  2. 方法二:

    python classify/train.py --model yolov5s-cls.pt --data datasets/data–epochs 100 --batch-size 32 --imgsz 224

训练完成后,会在runs文件夹下看到模型
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第10张图片

验证和推理

验证

  1. 方法一:

    python classify/val.py --weights runs/train-cls/exp4/weights/best.pt --data datasets/data

  2. 方法二:
    到classify文件夹下,修改val.py
    Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第11张图片
    第一行使用自己训练的模型文件,第二行修改数据集路径,其余参数可根据自行需要修改。

验证结果:
输出类别和正确率信息
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第12张图片

推理

方法与上述类似,该文件或直接使用命令行

# 测试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文件夹下查看结果

Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第13张图片
Yolov5 (v6.2) 使用自己的数据训练分类模型 基于ONNX TensorRT转换_第14张图片

导出

使用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

你可能感兴趣的:(人工智能,YOLO,分类,深度学习,onnx,tensorRT)