如何使用深度学习Yolov5来训练——树叶分类识别数据集 梧桐 银杏 樟叶 杜鹃 桂叶 广玉兰 木槿等10类 树叶分类识别检测

好的，以下是训练树叶分类识别YOLOv5模型的关键代码部分。我们将重点放在数据集加载、模型加载、训练配置和训练过程上。

1. 数据集配置文件 (data.yaml)

首先，确保你有一个正确的数据集配置文件 data.yaml，内容如下：

# data.yaml
train: Leaf-Classification-Dataset/images/train
val: Leaf-Classification-Dataset/images/val

nc: 10  # 类别数量
names: ['梧桐', '银杏', '樟叶', '杜鹃', '桂叶', '广玉兰', '木槿', '枫叶', '松叶', '竹叶']  # 类别名称

2. 训练脚本 (train_yolov5.py)

接下来是训练脚本的关键部分：

# train_yolov5.py
import torch
from yolov5.models.experimental import attempt_load
from yolov5.utils.datasets import create_dataloader
from yolov5.utils.general import check_img_size, increment_path
from yolov5.utils.torch_utils import select_device, time_synchronized
from yolov5.train import train

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device):
    # 选择设备
    device = select_device(device)

    # 加载预训练的YOLOv5模型
    model = attempt_load(model_config, map_location=device)

    # 设置数据集路径
    data_path = data_yaml_path

    # 开始训练
    train(
        data=data_path,
        epochs=epochs,  # 训练周期数
        batch_size=batch_size,  # 每批样本数量
        imgsz=img_size,  # 输入图像尺寸
        name="yolov5_leaf_classification",  # 输出模型的名字
        patience=10,  # 提早停止的耐心参数
        workers=4,  # 工作线程数
        device=device  # 设备（CPU或GPU）
    )

if __name__ == "__main__":
    data_yaml_path = 'data.yaml'
    model_config = 'yolov5s.pt'  # 你可以选择其他预训练模型，如'yolov5m.pt', 'yolov5l.pt'等
    epochs = 100
    batch_size = 16
    img_size = 640
    device = '0'  # 使用GPU，如果需要使用CPU，可以改为'cpu'

    train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device)

3. 关键代码解释

选择设备

device = select_device(device)

• select_device(device): 选择训练设备，可以是CPU或GPU。

加载预训练模型

model = attempt_load(model_config, map_location=device)

• attempt_load(model_config, map_location=device): 加载预训练的YOLOv5模型。

开始训练

train(
    data=data_path,
    epochs=epochs,  # 训练周期数
    batch_size=batch_size,  # 每批样本数量
    imgsz=img_size,  # 输入图像尺寸
    name="yolov5_leaf_classification",  # 输出模型的名字
    patience=10,  # 提早停止的耐心参数
    workers=4,  # 工作线程数
    device=device  # 设备（CPU或GPU）
)

• train(...): 调用YOLOv5的训练函数，传入训练配置参数。

4. 运行训练脚本

确保你的数据集路径和类别信息正确无误后，运行训练脚本：

python train_yolov5.py

5. 注意事项

1. 数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。
2. 模型配置：确保模型配置文件路径正确。
3. 图像大小：img_size可以根据实际需求调整，通常使用640或1280。
4. 设备：确保设备（CPU或GPU）可用。
5. 超参数调整：根据实际情况调整训练参数，如学习率、批量大小等，以获得最佳训练效果。