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

张

张建站

2026/8/10 4:24:47

10分钟阅读

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这类图形化标注工具安装LabelImgpip install labelImg labelImg # 启动标注工具标注流程打开图片文件夹设置标注格式为YOLO创建类别标签如defect_0, defect_1等绘制边界框并选择对应类别保存标注文件自动生成.txt文件标注注意事项确保边界框完全包含目标对于遮挡目标标注可见部分保持标注一致性同一类目标使用相同标注标准标注足够多的样本每个类别至少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.zip3. 开始模型训练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设备ID0表示第一块GPU --workers 8 \ # 数据加载线程数 --name custom_training \ # 实验名称 --project runs/train \ # 保存目录 --exist-ok \ # 允许覆盖已有实验 --patience 50 \ # 早停耐心值 --save-period 10 \ # 每10个epoch保存一次权重3.3 监控训练过程训练开始后YOLO-v5会自动启动TensorBoard你可以通过以下方式监控训练进度本地查看如果支持图形界面tensorboard --logdir runs/train查看训练日志训练过程中控制台会实时显示各项指标损失函数值box_loss, obj_loss, cls_loss精度指标precision, recall, mAP0.5, mAP0.5:0.95学习率变化当前epoch和剩余时间训练结果文件训练完成后在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)问题2GPU内存不足# 减小批次大小 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 304. 模型评估与优化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这个命令会输出详细的评估指标mAP0.5IoU阈值为0.5时的平均精度mAP0.5:0.95IoU阈值从0.5到0.95的平均精度精确率Precision预测为正的样本中真正为正的比例召回率Recall真正为正的样本中被预测为正的比例每个类别的AP每个类别的平均精度4.2 可视化检测结果查看模型在实际图片上的表现from PIL import Image import torch # 加载训练好的模型 model torch.hub.load(ultralytics/yolov5, custom, pathruns/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 # 默认是SGD3. 模型结构调整# 修改模型配置文件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 # 导出为CoreMLiOS/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_model5.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, pathruns/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(host0.0.0.0, port5000)5.4 集成到现有系统将YOLO-v5模型集成到你的应用程序中# 实时视频流处理示例 import cv2 import torch class YOLOv5Detector: def __init__(self, model_path, conf_threshold0.25, iou_threshold0.45): self.model torch.hub.load(ultralytics/yolov5, custom, pathmodel_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_pathNone): 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 --half6.4 模型剪枝与量化为了在边缘设备上部署可以对模型进行优化# 模型剪枝示例 import torch import torch.nn.utils.prune as prune def prune_model(model, amount0.3): 对模型进行剪枝 for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): # 对卷积层进行L1范数剪枝 prune.l1_unstructured(module, nameweight, amountamount) 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, inplaceTrue) # 校准需要校准数据集 # calibration_data ... # model(calibration_data) # 转换 torch.quantization.convert(model, inplaceTrue) return model6.5 持续学习与模型更新在实际应用中模型需要不断更新以适应新数据class ContinualLearning: def __init__(self, model_path): self.model torch.hub.load(ultralytics/yolov5, custom, pathmodel_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, epochs10): 在新数据上微调 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): 准备训练数据集 # 实现数据加载和预处理 # ... pass7. 总结通过本文的完整指南你应该已经掌握了YOLO-v5自定义训练的全流程。从数据准备、模型训练到部署应用每个环节都有其关键要点数据是基础高质量、多样化的标注数据是成功的第一步。确保标注准确、类别平衡、数据分布合理。训练要耐心不要期望一次训练就能得到完美模型。可能需要多次调整超参数、数据增强策略甚至重新审视数据质量。验证要全面不仅要看mAP指标还要在实际场景中测试模型表现。混淆矩阵、PR曲线、实际推理结果都是重要的评估依据。部署要考虑实际根据目标平台选择合适的模型格式和优化策略。边缘设备可能需要量化服务器端可能更关注吞吐量。持续改进模型上线后收集新的数据、监控模型表现、定期更新模型这是保持模型效果的关键。YOLO-v5的强大之处不仅在于其优秀的性能更在于其极致的工程友好性。它让自定义训练变得简单直接让更多开发者能够将目标检测技术应用到自己的业务场景中。记住最好的模型不是指标最高的模型而是最能解决实际问题的模型。开始你的自定义训练之旅吧让YOLO-v5为你的业务创造价值。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。