YOLOv5与PyTorch实战：Python物体检测推理全流程指南

一、技术背景与核心优势

YOLOv5作为Ultralytics团队开发的单阶段目标检测模型，凭借其速度与精度的平衡优势，已成为工业界和学术界的主流选择。PyTorch作为动态计算图框架，与YOLOv5的深度集成使得模型训练和部署更加灵活。本文将聚焦如何利用这两者构建完整的物体检测推理流程，适用于安防监控、自动驾驶、工业质检等场景。

1.1 YOLOv5模型特性

• 架构创新：基于CSPDarknet骨干网络，集成PANet特征融合模块，支持多尺度检测。
• 版本迭代：从v5s到v5x的4种规模模型，覆盖不同精度/速度需求（v5s-FP16推理可达140FPS）。
• 预训练权重：提供COCO数据集预训练模型，支持零代码迁移学习。

1.2 PyTorch生态优势

• 动态图机制：支持即时模式调试，便于模型结构修改。
• CUDA加速：自动利用GPU并行计算，推理速度较CPU提升10-50倍。
• TorchScript兼容：可将模型导出为中间表示，实现跨平台部署。

二、环境配置与依赖安装

2.1 系统要求

• Python 3.8+
• PyTorch 1.7+（推荐CUDA 11.x）
• CUDA 10.2+/cuDNN 8.0+（GPU环境）
• OpenCV 4.x（图像处理）

2.2 安装步骤

# 创建虚拟环境（推荐）
conda create -n yolov5_env python=3.8
conda activate yolov5_env
# 安装PyTorch（根据CUDA版本选择）
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
# 安装YOLOv5依赖
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

验证安装：

import torch
print(torch.__version__)  # 应输出1.7+
print(torch.cuda.is_available())  # GPU环境应返回True

三、模型加载与预处理

3.1 模型选择策略

模型版本	输入尺寸	mAP@0.5	推理速度（V100）	适用场景
yolov5s	640x640	56.8	140FPS	实时边缘设备
yolov5m	640x640	64.3	50FPS	通用嵌入式设备
yolov5l	640x640	67.3	30FPS	高精度监控系统
yolov5x	640x640	69.8	15FPS	云端离线分析

3.2 模型加载代码

from yolov5.models.experimental import attempt_load
import torch
# 加载预训练模型（自动下载）
model = attempt_load('yolov5s.pt', map_location='cuda' if torch.cuda.is_available() else 'cpu')
model.eval()  # 切换至推理模式

3.3 图像预处理流程

import cv2
import numpy as np
from yolov5.utils.augmentations import letterbox
def preprocess(img_path, img_size=640):
    # 读取图像
    img0 = cv2.imread(img_path)  # BGR格式
    assert img0 is not None, f'Image Not Found {img_path}'
    # 像素值归一化与通道转换
    img = letterbox(img0, img_size, stride=32, auto=True)[0]
    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to('cuda' if torch.cuda.is_available() else 'cpu')
    img = img.float() / 255.0  # 归一化到[0,1]
    if img.ndimension() == 3:
        img = img.unsqueeze(0)  # 添加batch维度
    return img0, img

四、推理执行与结果解析

4.1 核心推理代码

def detect(img_path, conf_thres=0.25, iou_thres=0.45):
    # 预处理
    img0, img = preprocess(img_path)
    # 推理（禁用梯度计算）
    with torch.no_grad():
        pred = model(img)[0]  # 输出包含检测结果
    # NMS后处理
    pred = non_max_suppression(pred, conf_thres, iou_thres)
    # 解析结果
    detections = []
    for det in pred:  # 每张图像的检测结果
        if len(det):
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()
            for *xyxy, conf, cls in reversed(det):
                label = f'{model.names[int(cls)]} {conf:.2f}'
                detections.append({
                    'bbox': [int(x) for x in xyxy],
                    'confidence': float(conf),
                    'class': model.names[int(cls)],
                    'label': label
                })
    return img0, detections

4.2 结果可视化实现

def plot_detections(img, detections):
    for det in detections:
        x1, y1, x2, y2 = det['bbox']
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = det['label']
        tf_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
        cv2.putText(img, label, (x1, y1 - tf_size[1] - 10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
    return img
# 使用示例
img_path = 'bus.jpg'
img_processed, detections = detect(img_path)
result_img = plot_detections(img_processed.copy(), detections)
cv2.imwrite('result.jpg', result_img)

五、性能优化与部署方案

5.1 推理速度优化

• 模型量化：使用TorchScript进行FP16量化，速度提升30%

  # 导出为TorchScript
  traced_script_module = torch.jit.trace(model, example_input)
  traced_script_module.save("yolov5s_quant.pt")

• TensorRT加速：NVIDIA GPU上可提升2-5倍吞吐量
• 多线程处理：使用concurrent.futures实现批量推理

5.2 部署方案对比

部署方式	适用场景	工具链	性能指标
PyTorch原生	研发调试阶段	torch.jit.trace	基准性能
ONNX Runtime	跨平台部署	ONNX转换 + ORT执行器	CPU推理加速30%
TensorRT	NVIDIA GPU生产环境	TRT引擎编译	延迟降低至2ms
TFLite	移动端/边缘设备	TFLite转换器	模型体积缩小4倍

六、常见问题解决方案

6.1 内存不足问题

• 现象：CUDA内存错误（RuntimeError: CUDA out of memory）
• 解决方案：
  • 降低img_size参数（如从640改为416）
  • 使用torch.cuda.empty_cache()清理缓存
  • 启用梯度检查点（model.half()进行混合精度）

6.2 检测精度下降

• 可能原因：
  • 输入图像分辨率与训练数据差异过大
  • 置信度阈值设置过高
  • 类别不平衡问题
• 优化建议：
  • 对特定场景进行微调训练
  • 调整conf_thres参数（默认0.25）
  • 使用WBF（Weighted Boxes Fusion）融合多尺度检测结果

七、进阶应用方向

7.1 自定义数据集训练

from yolov5.train import train
# 数据集结构要求
# datasets/
#   └── custom/
#       ├── images/
#       │   ├── train/
#       │   └── val/
#       └── labels/
#           ├── train/
#           └── val/
# 训练配置示例
data_dict = {
    'train': 'datasets/custom/images/train',
    'val': 'datasets/custom/images/val',
    'nc': 3,  # 类别数
    'names': ['class1', 'class2', 'class3']
}
train(data='custom.yaml', 
      weights='yolov5s.pt',
      img_size=640,
      batch_size=16,
      epochs=100)

7.2 视频流实时检测

def video_detection(source='0'):  # 0表示默认摄像头
    cap = cv2.VideoCapture(source)
    fps = cap.get(cv2.CAP_PROP_FPS)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        img, detections = detect(frame)
        result_frame = plot_detections(img, detections)
        cv2.imshow('YOLOv5 Detection', result_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()

八、总结与最佳实践

1. 模型选择原则：根据部署设备的计算能力选择适当规模的YOLOv5版本
2. 预处理标准化：保持与训练数据相同的归一化方式和尺寸调整策略
3. 后处理优化：合理设置NMS阈值（通常0.4-0.5）平衡精度与召回
4. 性能监控：使用torch.cuda.profiler分析GPU利用率
5. 持续更新：关注Ultralytics官方仓库的模型升级和bug修复

通过本文介绍的完整流程，开发者可以快速构建从图像输入到检测结果输出的端到端系统。实际应用中，建议结合具体场景进行模型微调和参数调优，以获得最佳检测效果。