如何使用深度学习框架训练使用——螺丝螺母目标检测数据集，训练测试共计400张，可用于数据分类，目标检测螺丝数据集 螺母数据集

如何使用这个螺丝螺母目标检测数据集进行YOLOv5模型的训练。我们将从数据集的准备、模型的加载、训练配置和训练过程等方面进行详细说明。

1. 数据集准备

数据集概述

• 数据集名称: 螺丝螺母目标检测数据集
• 数据集来源: 自制
• 数据集内容: 包含400张图像，每张图像都有对应的标签文件，标签文件采用YOLO格式。
• 检测目标: 包括螺丝和螺母。
• 数据集划分: 假设我们将其按8:1:1的比例划分为训练集、验证集和测试集，即320张训练集、40张验证集和40张测试集。
  

数据集目录结构

Screw-Nut-Dataset/
├── images/
│   ├── train/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── val/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   └── test/
│       ├── image1.jpg
│       ├── image2.jpg
│       └── ...
└── labels/
    ├── train/
    │   ├── image1.txt
    │   ├── image2.txt
    │   └── ...
    ├── val/
    │   ├── image1.txt
    │   ├── image2.txt
    │   └── ...
    └── test/
        ├── image1.txt
        ├── image2.txt
        └── ...

2. 数据集配置文件

创建一个data.yaml文件，配置数据集路径和类别信息。

# data.yaml
train: Screw-Nut-Dataset/images/train
val: Screw-Nut-Dataset/images/val
test: Screw-Nut-Dataset/images/test

nc: 2  # 类别数量
names: ['screw', 'nut']  # 类别名称

3. 划分数据集

如果你需要自己划分数据集，可以使用以下Python代码：

import os
import random
import shutil

# 数据集路径
dataset_path = 'Screw-Nut-Dataset'
images_path = os.path.join(dataset_path, 'images')
labels_path = os.path.join(dataset_path, 'labels')

# 创建目录
os.makedirs(os.path.join(images_path, 'train'), exist_ok=True)
os.makedirs(os.path.join(images_path, 'val'), exist_ok=True)
os.makedirs(os.path.join(images_path, 'test'), exist_ok=True)
os.makedirs(os.path.join(labels_path, 'train'), exist_ok=True)
os.makedirs(os.path.join(labels_path, 'val'), exist_ok=True)
os.makedirs(os.path.join(labels_path, 'test'), exist_ok=True)

# 获取所有图像和标签文件
all_images = [f for f in os.listdir(images_path) if f.endswith('.jpg')]
all_labels = [f for f in os.listdir(labels_path) if f.endswith('.txt')]

# 打乱顺序
random.shuffle(all_images)

# 划分数据集
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1

train_split = int(len(all_images) * train_ratio)
val_split = int(len(all_images) * (train_ratio + val_ratio))

train_images = all_images[:train_split]
val_images = all_images[train_split:val_split]
test_images = all_images[val_split:]

# 移动文件
for img in train_images:
    label = img.replace('.jpg', '.txt')
    shutil.move(os.path.join(images_path, img), os.path.join(images_path, 'train', img))
    shutil.move(os.path.join(labels_path, label), os.path.join(labels_path, 'train', label))

for img in val_images:
    label = img.replace('.jpg', '.txt')
    shutil.move(os.path.join(images_path, img), os.path.join(images_path, 'val', img))
    shutil.move(os.path.join(labels_path, label), os.path.join(labels_path, 'val', label))

for img in test_images:
    label = img.replace('.jpg', '.txt')
    shutil.move(os.path.join(images_path, img), os.path.join(images_path, 'test', img))
    shutil.move(os.path.join(labels_path, label), os.path.join(labels_path, 'test', label))

4. 训练脚本

创建一个训练脚本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_screw_nut",  # 输出模型的名字
        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)

5. 关键代码解释

选择设备

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_screw_nut",  # 输出模型的名字
    patience=10,  # 提早停止的耐心参数
    workers=4,  # 工作线程数
    device=device  # 设备（CPU或GPU）
)

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

6. 运行训练脚本

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

python train_yolov5.py

7. 注意事项

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

8. 总结

通过以上步骤，你可以使用YOLOv5训练一个针对螺丝螺母目标检测的高精度模型。