YOLOv5与PyTorch实战：从零实现高效物体检测

一、技术选型与核心优势

YOLOv5作为单阶段目标检测算法的代表，其核心优势在于速度与精度的平衡。基于PyTorch实现的版本（Ultralytics官方库）支持动态图模式，便于调试与自定义修改。相比TensorFlow版本，PyTorch的灵活性更适用于研究场景，且生态中预训练模型丰富（如COCO数据集训练的权重）。

关键特性：

• Anchor-Free设计：YOLOv5s起支持无锚框检测，减少超参数
• 自适应训练策略：自动调整学习率、批大小等参数
• 轻量化部署：模型可导出为ONNX/TensorRT格式，适配移动端

二、环境配置与依赖管理

2.1 基础环境搭建

推荐使用Anaconda创建隔离环境：

conda create -n yolov5_env python=3.8
conda activate yolov5_env
pip install torch torchvision torchaudio  # 根据CUDA版本选择版本号
pip install opencv-python matplotlib tqdm

2.2 YOLOv5官方库安装

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt  # 包含PyTorch及其他依赖

版本兼容性注意：

• PyTorch 1.7+ 推荐使用CUDA 11.x
• 若出现torch.cuda.is_available()为False，检查NVIDIA驱动与CUDA Toolkit安装

三、模型加载与基础推理

3.1 预训练模型使用

Ultralytics提供了多种规模的预训练模型：

import torch
from models.experimental import attempt_load
# 加载YOLOv5s模型（速度最快，精度适中）
model = attempt_load('yolov5s.pt', map_location='cuda')  # 或'cpu'
model.eval()  # 切换为推理模式
# 输入预处理（需归一化到[0,1]）
img = cv2.imread('test.jpg')[:, :, ::-1]  # BGR转RGB
img_tensor = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
img_tensor = img_tensor.unsqueeze(0).to('cuda')  # 添加batch维度
# 推理与后处理
with torch.no_grad():
    pred = model(img_tensor)
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)  # NMS处理

3.2 输出解析

预测结果为嵌套列表，每个元素对应一张图片的检测结果：

for det in pred:  # 遍历每张图片的检测结果
    if len(det):
        det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img.shape).round()  # 坐标还原
        for *xyxy, conf, cls in reversed(det):  # xyxy格式坐标
            label = f'{model.names[int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, img, label=label, color=(255,0,0))

四、自定义数据集训练

4.1 数据集准备

需符合YOLO格式：

dataset/
├── images/
│   ├── train/  # 训练图片
│   └── val/    # 验证图片
└── labels/
    ├── train/  # 对应.txt标签文件
    └── val/

标签文件格式（每行一个目标）：

class_id x_center y_center width height  # 均为相对值（0~1）

4.2 配置文件修改

创建data.yaml指定数据集路径与类别：

train: ../dataset/images/train
val: ../dataset/images/val
nc: 3  # 类别数
names: ['cat', 'dog', 'person']  # 类别名称

4.3 训练命令

python train.py --img 640 --batch 16 --epochs 100 \
                --data data.yaml --cfg models/yolov5s.yaml \
                --weights yolov5s.pt --name custom_model

关键参数说明：

• --img：输入分辨率（建议640或1280）
• --batch：根据GPU内存调整（V100可设64+）
• --epochs：通常300轮可达收敛
• --weights：''表示从头训练，yolov5s.pt表示迁移学习

五、模型优化技巧

5.1 数据增强策略

在data/augmentations.py中可自定义：

• Mosaic增强：4张图片拼接，提升小目标检测能力
• MixUp增强：图片混合，防止过拟合
• HSV空间调整：随机修改亮度、饱和度

5.2 超参数调优

学习率策略：

# 在train.py中修改optimizer
optimizer = torch.optim.SGD(model.parameters(), 
                           lr=0.01,  # 初始学习率
                           momentum=0.937,
                           weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.01, steps_per_epoch=len(train_loader),
    epochs=100, pct_start=0.1)  # 10%周期线性增长学习率

5.3 模型剪枝与量化

通道剪枝示例：

from models.yolo import prune_model
# 保留70%的通道
pruned_model = prune_model(model, ratio=0.3)
torch.save(pruned_model.state_dict(), 'pruned.pt')

六、部署与性能优化

6.1 ONNX导出

img_size = 640
dummy_input = torch.randn(1, 3, img_size, img_size).to('cuda')
torch.onnx.export(model,
                 dummy_input,
                 'yolov5s.onnx',
                 input_names=['images'],
                 output_names=['output'],
                 dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}},
                 opset_version=12)

6.2 TensorRT加速（NVIDIA GPU）

trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine --fp16

性能对比（以Tesla T4为例）：
| 格式 | 延迟(ms) | 吞吐量(FPS) |
|——————|—————|——————-|
| PyTorch | 12.5 | 80 |
| ONNX | 9.8 | 102 |
| TensorRT | 6.2 | 161 |

七、常见问题解决方案

7.1 CUDA内存不足

• 减小--batch-size（如从16降至8）
• 使用梯度累积：

  accumulate = 4  # 每4个batch更新一次权重
  for i, (imgs, targets) in enumerate(dataloader):
    loss = model(imgs, targets)
    loss.backward()
    if (i + 1) % accumulate == 0:
        optimizer.step()
        optimizer.zero_grad()

7.2 检测框抖动

• 增加NMS阈值（--iou-thres 0.5→0.6）
• 使用指数移动平均（EMA）模型权重：

  # 在train.py中启用
  ema = ModelEMA(model, 0.9998)  # 0.9998是常用衰减率
  ...
  ema.update(model)  # 每个step后调用

八、进阶研究方向

1. 多尺度检测改进：在FPN结构中增加特征层
2. 注意力机制融合：在YOLO头中加入CBAM模块
3. 半监督学习：利用未标注数据通过Pseudo Labeling训练

实践建议：

• 从YOLOv5s开始实验，逐步增加模型复杂度
• 使用Weights & Biases等工具记录实验过程
• 关注Ultralytics官方仓库的更新（如YOLOv8的改进）

通过本文的完整流程，开发者可快速掌握从数据准备到部署的全链条技术，实现高效的物体检测系统。实际项目中，建议结合具体场景调整模型规模与后处理阈值，以达到速度与精度的最佳平衡。