当前位置： 首页 > news >正文

YOLO-v5自定义训练：在自己的数据集上微调模型

news 2026/3/27 3:40:37

YOLO-v5自定义训练：在自己的数据集上微调模型

如果你正在寻找一个既快又准的目标检测解决方案，YOLO-v5很可能已经进入了你的视野。这个没有发表过顶会论文的模型，却凭借其出色的工程化设计和极致的易用性，成为了工业界部署最广的目标检测框架之一。

但你可能面临一个实际问题：预训练模型在COCO数据集上表现再好，也无法直接识别你的特定目标——无论是生产线上的瑕疵零件、医疗影像中的病灶区域，还是农业无人机拍摄的病虫害叶片。

这就是自定义训练的价值所在。本文将带你一步步完成YOLO-v5在自己数据集上的完整微调流程，从数据准备到模型部署，让你真正掌握这个强大工具的使用方法。

1. 环境准备与快速部署

1.1 获取YOLO-v5镜像

在CSDN星图平台上，你可以直接找到预配置好的YOLO-v5镜像。这个镜像已经包含了PyTorch深度学习框架、YOLO-v5专用工具库以及完整的开发环境，省去了繁琐的环境配置步骤。

选择YOLO-v5镜像后，系统会自动为你创建一个包含所有必要依赖的容器环境。启动后，你可以通过Jupyter Notebook或SSH两种方式访问：

  • Jupyter方式：通过Web界面直接访问，适合交互式开发和调试
  • SSH方式：通过命令行访问，适合批量处理和自动化任务

1.2 验证环境配置

环境启动后，首先验证YOLO-v5是否正常工作。进入项目目录并运行一个简单的测试：

cd /root/yolov5/ python -c "import torch; print('PyTorch版本:', torch.__version__)"

接着测试YOLO-v5的基本功能：

import torch # 加载预训练的YOLO-v5模型（可选模型：yolov5n, yolov5s, yolov5m, yolov5l, yolov5x） model = torch.hub.load("ultralytics/yolov5", "yolov5s") # 使用示例图片进行推理 img = "https://ultralytics.com/images/zidane.jpg" results = model(img) # 查看结果 results.print() # 打印检测结果到控制台 results.show() # 显示检测结果（需要图形界面）

如果一切正常，你会看到模型成功检测出图片中的人物和物体，并输出相应的边界框和置信度。

2. 准备自定义数据集

2.1 数据标注格式

YOLO-v5使用特定的标注格式，每个标注文件对应一张图片，包含该图片中所有目标的标注信息。标注文件的格式如下：

<类别ID> <中心点x坐标> <中心点y坐标> <宽度> <高度>

所有坐标值都是相对于图片宽度和高度的归一化值（0-1之间）。例如，如果一张图片尺寸为640x480，一个目标的边界框左上角在(100, 50)，右下角在(300, 200)，那么标注应该是：

0 0.3125 0.2604 0.3125 0.3125

其中：

  • 0是类别ID（需要与你的类别对应）
  • 0.3125 = (100 + 300)/2 / 640（中心点x坐标）
  • 0.2604 = (50 + 200)/2 / 480（中心点y坐标）
  • 0.3125 = (300 - 100) / 640（宽度）
  • 0.3125 = (200 - 50) / 480（高度）

2.2 使用标注工具

对于大多数用户，我推荐使用LabelImg或CVAT这类图形化标注工具：

  1. 安装LabelImg：
pip install labelImg labelImg # 启动标注工具

  1. 标注流程：

    • 打开图片文件夹
    • 设置标注格式为YOLO
    • 创建类别标签（如：defect_0, defect_1等）
    • 绘制边界框并选择对应类别
    • 保存标注文件（自动生成.txt文件）

  2. 标注注意事项：

    • 确保边界框完全包含目标
    • 对于遮挡目标，标注可见部分
    • 保持标注一致性（同一类目标使用相同标注标准）
    • 标注足够多的样本（每个类别至少100-200张）

2.3 组织数据集结构

YOLO-v5要求特定的目录结构。创建一个标准的数据集目录如下：

custom_dataset/ ├── images/ │ ├── train/ # 训练图片 │ │ ├── image1.jpg │ │ ├── image2.jpg │ │ └── ... │ └── val/ # 验证图片 │ ├── image101.jpg │ ├── image102.jpg │ └── ... └── labels/ ├── train/ # 训练标注 │ ├── image1.txt │ ├── image2.txt │ └── ... └── val/ # 验证标注 ├── image101.txt ├── image102.txt └── ...

2.4 创建数据集配置文件

在YOLO-v5目录下创建数据集配置文件data/custom.yaml：

# 数据集路径（相对路径或绝对路径） path: ../custom_dataset # 数据集根目录 train: images/train # 训练集图片路径 val: images/val # 验证集图片路径 # 类别数量 nc: 3 # 你的数据集中有多少个类别 # 类别名称列表 names: ['cat', 'dog', 'person'] # 替换为你的实际类别名称 # 可选：下载地址/自动下载设置 # download: https://ultralytics.com/assets/coco128.zip

3. 开始模型训练

3.1 基础训练命令

准备好数据集后，就可以开始训练了。最基本的训练命令如下：

cd /root/yolov5/ python train.py --img 640 --batch 16 --epochs 100 --data data/custom.yaml --weights yolov5s.pt

这个命令会：

  • 使用640x640的输入尺寸
  • 设置批次大小为16（根据GPU内存调整）
  • 训练100个epoch
  • 使用你的自定义数据集配置
  • 从YOLOv5s预训练权重开始微调

3.2 关键训练参数详解

YOLO-v5提供了丰富的训练参数，以下是一些最常用的：

python train.py \ --weights yolov5s.pt \ # 预训练权重 --data data/custom.yaml \ # 数据集配置文件 --epochs 300 \ # 训练轮数 --batch-size 16 \ # 批次大小 --img-size 640 \ # 输入图片尺寸 --device 0 \ # GPU设备ID（0表示第一块GPU） --workers 8 \ # 数据加载线程数 --name custom_training \ # 实验名称 --project runs/train \ # 保存目录 --exist-ok \ # 允许覆盖已有实验 --patience 50 \ # 早停耐心值 --save-period 10 \ # 每10个epoch保存一次权重

3.3 监控训练过程

训练开始后，YOLO-v5会自动启动TensorBoard，你可以通过以下方式监控训练进度：

  1. 本地查看（如果支持图形界面）：
tensorboard --logdir runs/train
  1. 查看训练日志： 训练过程中，控制台会实时显示各项指标：
  • 损失函数值（box_loss, obj_loss, cls_loss）
  • 精度指标（precision, recall, mAP@0.5, mAP@0.5:0.95）
  • 学习率变化
  • 当前epoch和剩余时间
  1. 训练结果文件： 训练完成后，在runs/train/exp目录下会生成：
  • weights/：最佳权重和最后权重
  • results.png：训练指标曲线图
  • confusion_matrix.png：混淆矩阵
  • labels.jpg：标注可视化
  • opt.yaml：训练配置

3.4 解决常见训练问题

在训练自定义数据集时，你可能会遇到以下问题：

问题1：损失不下降或波动很大

# 尝试降低学习率 python train.py --lr 0.01 # 默认是0.1 # 增加热身epoch python train.py --warmup-epochs 5 # 默认是3 # 检查数据标注质量 python -c "from utils.general import check_dataset; check_dataset('data/custom.yaml')"

问题2：GPU内存不足

# 减小批次大小 python train.py --batch-size 8 # 减小输入尺寸 python train.py --img 416 # 使用更小的模型 python train.py --weights yolov5n.pt

问题3：过拟合（训练集好，验证集差）

# 增加数据增强 python train.py --augment True # 使用权重衰减 python train.py --weight-decay 0.0005 # 早停策略 python train.py --patience 30

4. 模型评估与优化

4.1 评估训练结果

训练完成后，使用验证集评估模型性能：

python val.py \ --weights runs/train/exp/weights/best.pt \ --data data/custom.yaml \ --img 640 \ --batch-size 32 \ --task val \ --name custom_eval \ --save-json \ --save-txt \ --save-conf

这个命令会输出详细的评估指标：

  • mAP@0.5：IoU阈值为0.5时的平均精度
  • mAP@0.5:0.95：IoU阈值从0.5到0.95的平均精度
  • 精确率（Precision）：预测为正的样本中真正为正的比例
  • 召回率（Recall）：真正为正的样本中被预测为正的比例
  • 每个类别的AP：每个类别的平均精度

4.2 可视化检测结果

查看模型在实际图片上的表现：

from PIL import Image import torch # 加载训练好的模型 model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt') # 测试单张图片 img_path = 'path/to/your/test_image.jpg' results = model(img_path) # 显示结果 results.show() # 保存结果 results.save('output/') # 获取检测结果的详细信息 predictions = results.pandas().xyxy[0] print(predictions)

4.3 模型性能分析

使用YOLO-v5内置的分析工具深入了解模型表现：

from utils.plots import plot_results import matplotlib.pyplot as plt # 绘制训练曲线 plot_results('runs/train/exp/results.csv') # 保存为图片 # 分析混淆矩阵 from utils.plots import plot_confusion_matrix plot_confusion_matrix('runs/train/exp/confusion_matrix.png') # 计算每个类别的性能 from utils.metrics import ap_per_class import numpy as np # 这里需要你的预测结果和真实标签 # predictions, targets = ... # ap, p, r = ap_per_class(predictions, targets)

4.4 模型优化技巧

如果你的模型表现不够理想，可以尝试以下优化方法：

1. 数据层面优化

# 检查数据平衡性 import os from collections import Counter label_dir = 'custom_dataset/labels/train/' class_counts = Counter() for label_file in os.listdir(label_dir): with open(os.path.join(label_dir, label_file), 'r') as f: for line in f: class_id = int(line.split()[0]) class_counts[class_id] += 1 print("各类别样本数量:", class_counts)

2. 超参数调优

# 尝试不同的学习率策略 python train.py --lr0 0.01 --lrf 0.1 # 初始学习率0.01，最终学习率0.001 # 调整数据增强强度 python train.py --hsv-h 0.015 --hsv-s 0.7 --hsv-v 0.4 # HSV增强参数 python train.py --degrees 10.0 --translate 0.1 --scale 0.5 # 空间变换参数 # 使用不同的优化器 python train.py --optimizer AdamW # 默认是SGD

3. 模型结构调整

# 修改模型配置文件（models/yolov5s.yaml） # 调整neck或head的结构 # 注意：这需要一定的模型架构知识

5. 模型部署与应用

5.1 导出为不同格式

YOLO-v5支持导出多种格式，方便在不同平台上部署：

# 导出为TorchScript python export.py --weights runs/train/exp/weights/best.pt --include torchscript # 导出为ONNX（推荐） python export.py --weights runs/train/exp/weights/best.pt --include onnx # 导出为TensorRT（需要CUDA） python export.py --weights runs/train/exp/weights/best.pt --include engine --device 0 # 导出为CoreML（iOS/macOS） python export.py --weights runs/train/exp/weights/best.pt --include coreml # 导出为TensorFlow SavedModel python export.py --weights runs/train/exp/weights/best.pt --include saved_model

5.2 使用导出的模型进行推理

导出的模型可以直接用于推理，无需原始训练代码：

# 使用ONNX模型推理 import onnxruntime import cv2 import numpy as np # 加载ONNX模型 session = onnxruntime.InferenceSession('best.onnx') # 预处理图片 img = cv2.imread('test.jpg') img = cv2.resize(img, (640, 640)) img = img.transpose(2, 0, 1) # HWC to CHW img = img[np.newaxis, ...] # 添加batch维度 img = img.astype(np.float32) / 255.0 # 归一化 # 推理 inputs = {session.get_inputs()[0].name: img} outputs = session.run(None, inputs) # 后处理（需要根据模型输出格式调整） # ...

5.3 创建简单的Web服务

将训练好的模型封装为API服务：

# app.py from flask import Flask, request, jsonify import torch from PIL import Image import io app = Flask(__name__) # 加载模型 model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt') @app.route('/predict', methods=['POST']) def predict(): if 'file' not in request.files: return jsonify({'error': 'No file provided'}), 400 file = request.files['file'] img_bytes = file.read() img = Image.open(io.BytesIO(img_bytes)) # 推理 results = model(img) # 解析结果 predictions = results.pandas().xyxy[0].to_dict('records') return jsonify({ 'predictions': predictions, 'image_size': img.size }) if __name__ == '__main__': app.run(host='0.0.0.0', port=5000)

5.4 集成到现有系统

将YOLO-v5模型集成到你的应用程序中：

# 实时视频流处理示例 import cv2 import torch class YOLOv5Detector: def __init__(self, model_path, conf_threshold=0.25, iou_threshold=0.45): self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path) self.model.conf = conf_threshold # 置信度阈值 self.model.iou = iou_threshold # IoU阈值 def process_frame(self, frame): # 转换BGR到RGB rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 推理 results = self.model(rgb_frame) # 绘制检测框 annotated_frame = results.render()[0] # 转换回BGR annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR) return annotated_frame, results.pandas().xyxy[0] def process_video(self, video_path, output_path=None): cap = cv2.VideoCapture(video_path) if output_path: fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4)))) while cap.isOpened(): ret, frame = cap.read() if not ret: break processed_frame, detections = self.process_frame(frame) if output_path: out.write(processed_frame) cv2.imshow('Detection', processed_frame) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() if output_path: out.release() cv2.destroyAllWindows() # 使用示例 detector = YOLOv5Detector('runs/train/exp/weights/best.pt') detector.process_video('input_video.mp4', 'output_video.mp4')

6. 进阶技巧与最佳实践

6.1 使用预训练权重的技巧

YOLO-v5提供了多种预训练权重，选择合适的起点很重要：

# 不同规模的预训练模型 yolov5n.pt # 纳米版：最小最快，适合移动端 yolov5s.pt # 小版：平衡速度与精度 yolov5m.pt # 中版：更好的精度 yolov5l.pt # 大版：高精度 yolov5x.pt # 超大版：最高精度 # 根据你的需求选择： # 1. 如果数据与COCO类似 → 直接用对应模型 # 2. 如果数据差异大 → 从yolov5s开始 # 3. 如果计算资源有限 → 从yolov5n开始 # 4. 如果需要最高精度 → 从yolov5x开始

6.2 数据增强策略调整

YOLO-v5内置了强大的数据增强，但你可以根据具体任务调整：

# 在data/custom.yaml中添加增强配置 augment: true augmentation: hsv_h: 0.015 # 色调增强（0-1） hsv_s: 0.7 # 饱和度增强（0-1） hsv_v: 0.4 # 明度增强（0-1） degrees: 0.0 # 旋转角度（±度） translate: 0.1 # 平移（±比例） scale: 0.5 # 缩放（±比例） shear: 0.0 # 剪切（±度） perspective: 0.0 # 透视变换（0-0.001） flipud: 0.0 # 上下翻转概率 fliplr: 0.5 # 左右翻转概率 mosaic: 1.0 # mosaic增强概率 mixup: 0.0 # mixup增强概率

6.3 多GPU训练与分布式训练

对于大数据集，可以使用多GPU加速训练：

# 使用多GPU训练（数据并行） python train.py --device 0,1,2,3 # 使用4块GPU # 使用分布式训练 python -m torch.distributed.launch --nproc_per_node 4 train.py \ --weights yolov5s.pt \ --data custom.yaml \ --epochs 300 \ --batch-size 64 # 总批次大小，每GPU 16 # 使用混合精度训练（减少显存，加快训练） python train.py --device 0 --batch-size 32 --half

6.4 模型剪枝与量化

为了在边缘设备上部署，可以对模型进行优化：

# 模型剪枝示例 import torch import torch.nn.utils.prune as prune def prune_model(model, amount=0.3): """对模型进行剪枝""" for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): # 对卷积层进行L1范数剪枝 prune.l1_unstructured(module, name='weight', amount=amount) prune.remove(module, 'weight') return model # 模型量化（INT8） def quantize_model(model_path): """将模型量化为INT8""" model = torch.load(model_path) model.eval() # 准备量化配置 model.qconfig = torch.quantization.get_default_qconfig('fbgemm') # 准备量化 torch.quantization.prepare(model, inplace=True) # 校准（需要校准数据集） # calibration_data = ... # model(calibration_data) # 转换 torch.quantization.convert(model, inplace=True) return model

6.5 持续学习与模型更新

在实际应用中，模型需要不断更新以适应新数据：

class ContinualLearning: def __init__(self, model_path): self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path) self.new_data = [] def add_new_samples(self, images, labels): """添加新样本到记忆库""" self.new_data.extend(zip(images, labels)) # 如果样本太多，进行采样 if len(self.new_data) > 1000: self.new_data = self.new_data[-1000:] def fine_tune(self, epochs=10): """在新数据上微调""" if not self.new_data: return # 准备新数据集 new_dataset = self._prepare_dataset(self.new_data) # 微调训练 self.model.train() # 这里简化了训练过程，实际需要完整的训练循环 # ... print(f"在新数据上微调了{epochs}个epoch") def _prepare_dataset(self, data): """准备训练数据集""" # 实现数据加载和预处理 # ... pass

7. 总结

通过本文的完整指南，你应该已经掌握了YOLO-v5自定义训练的全流程。从数据准备、模型训练到部署应用，每个环节都有其关键要点：

数据是基础：高质量、多样化的标注数据是成功的第一步。确保标注准确、类别平衡、数据分布合理。

训练要耐心：不要期望一次训练就能得到完美模型。可能需要多次调整超参数、数据增强策略，甚至重新审视数据质量。

验证要全面：不仅要看mAP指标，还要在实际场景中测试模型表现。混淆矩阵、PR曲线、实际推理结果都是重要的评估依据。

部署要考虑实际：根据目标平台选择合适的模型格式和优化策略。边缘设备可能需要量化，服务器端可能更关注吞吐量。

持续改进：模型上线后，收集新的数据、监控模型表现、定期更新模型，这是保持模型效果的关键。

YOLO-v5的强大之处不仅在于其优秀的性能，更在于其极致的工程友好性。它让自定义训练变得简单直接，让更多开发者能够将目标检测技术应用到自己的业务场景中。

记住，最好的模型不是指标最高的模型，而是最能解决实际问题的模型。开始你的自定义训练之旅吧，让YOLO-v5为你的业务创造价值。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。

http://www.jsqmd.com/news/495770/

相关文章：

  • 一键部署DeerFlow镜像：火山引擎FaaS应用中心快速体验AI研究助理
  • 开发者必看：CosyVoice-300M Lite镜像部署实操手册，开箱即用
  • 黄山派小智动态待机界面进阶：从GIF优化到性能调优
  • VSCode 2026日志插件深度评测：性能提升273%、错误定位提速8.6倍，实测数据全公开
  • Docker容器间通信的3种实用方法：从host.docker.internal到自定义网络
  • Doris在大数据处理中的性能优化秘籍
  • Vue3项目实战：vue-cropper图片裁剪从安装到跨域问题全解决
  • 智谱开源视觉大模型GLM-4.6V-Flash-WEB体验：部署简单，响应快，效果惊艳
  • 微信小程序订阅消息授权数据的后端存储机制解析
  • GDSDecomp全解析：Godot游戏逆向工程实战指南
  • 计算机毕业设计java基于微信小程序的菜谱查询点评系统设计与开发 基于微信小程序的美食菜谱分享与评价系统 基于微信小程序的食谱查询与用户点评平台
  • Packet Tracer实验复盘：配置完RIP路由后，别忘了用这几个命令验证和排错
  • Qwen3-ASR-1.7B在媒体行业的应用：采访录音自动转写系统
  • el-cascader远程搜索避坑指南：从filterable到lazy加载的完整配置
  • 解决MTK手机自动亮度太亮/太暗问题：手动调整config.xml的完整流程
  • 从零开始：使用Docker容器化部署Django项目到腾讯云CVM（附完整配置文件）
  • 深入解析Chrome CORS跨域限制及实战解决方案
  • 基于强化学习的图片旋转判断模型优化
  • Harmonyos应用实例119：立体图形展开折叠游戏
  • 国企工程建筑局域网Web应用如何基于JS实现BIM模型大文件的目录结构分片传输？
  • 揭开Deliberate模型的实战面纱：从技术优势到落地挑战
  • 新手避坑指南：Simulink中AC Voltage Source最常见的3个配置错误及解决方法
  • 无需代码！Z-Image-Turbo_UI界面小白教程：开箱即用，快速出图
  • CLIP-GmP-ViT-L-14图文匹配测试工具：辅助数据库课程设计中的多媒体检索模块
  • 惊艳！bge-large-zh-v1.5向量生成效果展示：sglang部署实测案例
  • AI编程新时代：Mirage Flow实战自动化代码审查与重构
  • 企业数字化转型的核心策略与实践路径 - 企业推荐官【官方】
  • 医疗局域网如何通过百度WebUploader组件优化病历PDF文件的浏览器端分片断点恢复？
  • Overleaf本地化部署实战：从零搭建高效LaTeX协作环境
  • 深入解析Techpoint TP2855视频解码芯片的寄存器配置与应用（第四部分）