从Kaggle猫狗数据到可运行模型：YOLOv5s.pt训练全流程拆解（附完整代码与配置文件）

从Kaggle猫狗数据到可运行模型YOLOv5s.pt训练全流程拆解在计算机视觉领域目标检测一直是一个热门研究方向。YOLO(You Only Look Once)系列算法因其速度快、精度高的特点成为工业界和学术界广泛使用的目标检测框架。本文将详细介绍如何使用YOLOv5s.pt模型在Kaggle猫狗数据集上完成端到端的训练流程包括数据准备、环境配置、模型训练和结果验证等关键步骤。1. 环境准备与工具安装1.1 硬件与基础软件要求YOLOv5训练对硬件有一定要求建议配置GPUNVIDIA显卡推荐RTX 2060及以上显存≥6GB内存≥16GB存储空间至少50GB可用空间用于存放数据集和模型操作系统Windows 10/11或LinuxUbuntu 18.04基础软件依赖# 检查NVIDIA驱动是否安装 nvidia-smi # 输出应显示GPU信息如CUDA版本等1.2 Python环境配置推荐使用conda创建独立Python环境conda create -n yolov5 python3.8 conda activate yolov5安装PyTorch注意选择与CUDA版本匹配的安装命令# 以CUDA 11.3为例 pip install torch1.12.1cu113 torchvision0.13.1cu113 torchaudio0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113验证PyTorch是否正确识别GPUimport torch print(torch.__version__) # 应显示1.12.1 print(torch.cuda.is_available()) # 应输出True2. 数据集准备与标注2.1 获取Kaggle猫狗数据集访问Kaggle官网下载Dogs vs Cats数据集解压后目录结构应如下dogs-vs-cats/ ├── train/ │ ├── cat.0.jpg │ ├── dog.0.jpg │ └── ... └── test/ ├── 1.jpg └── ...2.2 数据标注工具使用推荐使用labelImg进行标注pip install labelImg labelImg # 启动标注工具标注时注意选择YOLO格式输出类别命名规范cat和dog保存的txt文件每行格式class_id x_center y_center width height2.3 数据集组织与划分最终数据集目录结构custom_data/ ├── images/ │ ├── train/ # 训练图片 │ └── val/ # 验证图片 └── labels/ ├── train/ # 训练标签 └── val/ # 验证标签可使用以下Python代码划分训练集和验证集import os import random from shutil import copyfile def split_dataset(src_img, src_label, train_ratio0.8): # 创建目录结构 os.makedirs(custom_data/images/train, exist_okTrue) os.makedirs(custom_data/images/val, exist_okTrue) os.makedirs(custom_data/labels/train, exist_okTrue) os.makedirs(custom_data/labels/val, exist_okTrue) # 获取所有图片文件 files [f for f in os.listdir(src_img) if f.endswith(.jpg)] random.shuffle(files) # 划分数据集 split_idx int(len(files)*train_ratio) for i, file in enumerate(files): base os.path.splitext(file)[0] if i split_idx: # 训练集 copyfile(f{src_img}/{file}, fcustom_data/images/train/{file}) copyfile(f{src_label}/{base}.txt, fcustom_data/labels/train/{base}.txt) else: # 验证集 copyfile(f{src_img}/{file}, fcustom_data/images/val/{file}) copyfile(f{src_label}/{base}.txt, fcustom_data/labels/val/{base}.txt)3. YOLOv5模型配置3.1 下载YOLOv5代码和预训练模型git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt下载预训练模型wget https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt -P weights/3.2 配置文件修改数据配置文件(data/custom.yaml):# 训练和验证数据路径 train: ../custom_data/images/train val: ../custom_data/images/val # 类别数量和名称 nc: 2 names: [cat, dog]模型配置文件(models/yolov5s_custom.yaml):# 参数 nc: 2 # 类别数 depth_multiple: 0.33 # 模型深度倍数 width_multiple: 0.50 # 层通道倍数 # 锚点框 anchors: - [10,13, 16,30, 33,23] # P3/8 - [30,61, 62,45, 59,119] # P4/16 - [116,90, 156,198, 373,326] # P5/32 # 主干网络 backbone: # [from, number, module, args] [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 [-1, 3, C3, [128]], [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 [-1, 6, C3, [256]], [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 [-1, 9, C3, [512]], [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 [-1, 3, C3, [1024]], [-1, 1, SPPF, [1024, 5]], # 9 ] # 检测头 head: [[-1, 1, Conv, [512, 1, 1]], [-1, 1, nn.Upsample, [None, 2, nearest]], [[-1, 6], 1, Concat, [1]], # cat backbone P4 [-1, 3, C3, [512, False]], # 13 [-1, 1, Conv, [256, 1, 1]], [-1, 1, nn.Upsample, [None, 2, nearest]], [[-1, 4], 1, Concat, [1]], # cat backbone P3 [-1, 3, C3, [256, False]], # 17 (P3/8-small) [-1, 1, Conv, [256, 3, 2]], [[-1, 14], 1, Concat, [1]], # cat head P4 [-1, 3, C3, [512, False]], # 20 (P4/16-medium) [-1, 1, Conv, [512, 3, 2]], [[-1, 10], 1, Concat, [1]], # cat head P5 [-1, 3, C3, [1024, False]], # 23 (P5/32-large) [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) ]4. 模型训练与调优4.1 启动训练基本训练命令python train.py --img 640 --batch 16 --epochs 100 --data data/custom.yaml --cfg models/yolov5s_custom.yaml --weights weights/yolov5s.pt --name cat_dog_detection关键参数说明参数说明推荐值--img输入图像尺寸640--batch批次大小根据GPU显存调整--epochs训练轮数100-300--data数据配置文件data/custom.yaml--cfg模型配置文件models/yolov5s_custom.yaml--weights预训练权重weights/yolov5s.pt--name实验名称自定义4.2 训练监控与可视化YOLOv5会自动记录训练日志可通过TensorBoard查看tensorboard --logdir runs/train主要监控指标损失函数box_loss, obj_loss, cls_loss评估指标precision, recall, mAP0.5, mAP0.5:0.95学习率当前学习率变化4.3 训练调优技巧学习率调整初始学习率0.01使用余弦退火调度器可通过--lr0参数调整数据增强默认启用Mosaic增强可通过--hyp指定超参数文件调整增强强度早停机制设置--patience参数如50当指标不再提升时自动停止训练5. 模型验证与测试5.1 验证集评估python val.py --weights runs/train/cat_dog_detection/weights/best.pt --data data/custom.yaml --img 640输出示例Class Images Instances P R mAP.5 mAP.5:.95 all 100 1500 0.95 0.92 0.94 0.72 cat 100 750 0.96 0.93 0.95 0.73 dog 100 750 0.94 0.91 0.93 0.715.2 单张图片测试python detect.py --weights runs/train/cat_dog_detection/weights/best.pt --source test_image.jpg --conf 0.5参数说明--source可指定图片、视频或摄像头0--conf置信度阈值--save-txt保存检测结果为YOLO格式5.3 性能测试python test.py --weights runs/train/cat_dog_detection/weights/best.pt --data data/custom.yaml --task speed输出示例Fusing layers... Model summary: 224 layers, 7266973 parameters, 0 gradients Speed: 2.0ms pre-process, 4.5ms inference, 1.0ms NMS per image at shape (1, 3, 640, 640)6. 常见问题与解决方案6.1 显存不足问题现象训练时出现CUDA out of memory错误解决方案减小--batch-size如从16降到8减小--img-size如从640降到416使用更小的模型如yolov5n.pt6.2 类别不平衡问题现象某个类别如猫的检测效果明显差于另一类别解决方案检查数据集标注是否正确使用类别权重修改损失函数增加少数类别的数据增强6.3 模型过拟合问题现象训练集指标很好但验证集指标差解决方案增加数据增强强度使用早停机制添加正则化如Dropout减少模型复杂度7. 模型部署与应用7.1 模型导出为ONNX格式python export.py --weights runs/train/cat_dog_detection/weights/best.pt --include onnx --img 6407.2 使用OpenCV调用模型import cv2 import numpy as np # 加载模型 net cv2.dnn.readNet(best.onnx) # 预处理图像 image cv2.imread(test.jpg) blob cv2.dnn.blobFromImage(image, 1/255.0, (640, 640), swapRBTrue, cropFalse) # 推理 net.setInput(blob) outputs net.forward(net.getUnconnectedOutLayersNames()) # 后处理 # ... (解析输出结果)7.3 模型量化与优化FP16量化python export.py --weights best.pt --include onnx --halfTensorRT加速python export.py --weights best.pt --include engine --device 0在实际项目中我们通常会根据硬件平台选择最优的部署方式。例如在边缘设备上可以使用TensorRT加速而在服务器端可以使用ONNX Runtime进行推理。