logo

开发者热搜

基于YOLOv5与PyTorch的Python物体检测推理全流程指南

作者：Nicky2025.10.12 02:44浏览量：6

简介：本文详细阐述如何使用YOLOv5模型与PyTorch框架在Python环境中实现高效的物体检测推理，涵盖环境配置、模型加载、预处理、推理执行及结果可视化全流程，适合开发者快速部署应用。

基于YOLOv5与PyTorch的Python物体检测推理全流程指南

一、技术背景与核心价值

物体检测是计算机视觉的核心任务之一，广泛应用于安防监控、自动驾驶、工业质检等领域。YOLOv5作为Ultralytics团队开发的实时目标检测框架，凭借其轻量化设计、高精度和易用性，成为工业界和学术界的热门选择。结合PyTorch的动态计算图特性与Python的生态优势，开发者可快速构建从训练到部署的全流程解决方案。本文将系统解析如何利用YOLOv5与PyTorch在Python中实现高效的物体检测推理，重点覆盖环境配置、模型加载、数据预处理、推理执行及结果可视化等关键环节。

二、环境配置：构建开发基础

1. Python环境要求

  • 版本选择：推荐Python 3.8+（YOLOv5官方测试环境）
  • 虚拟环境管理：使用conda create -n yolo_env python=3.8创建独立环境，避免依赖冲突
  • PyTorch安装：根据CUDA版本选择对应版本（如pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113）

2. YOLOv5安装方式

  • 源码安装（推荐）：
    1. git clone https://github.com/ultralytics/yolov5.git
    2. cd yolov5
    3. pip install -r requirements.txt
  • PyPI安装：pip install yolov5（适用于快速原型开发）

3. 验证环境

执行以下命令验证安装：

  1. import torch
  2. from yolov5 import detect
  3. print(f"PyTorch版本: {torch.__version__}")
  4. print(f"CUDA可用: {torch.cuda.is_available()}")

三、模型加载与配置

1. 预训练模型选择

YOLOv5提供多种规模模型：

  • YOLOv5s：轻量级（7.3M参数），适合移动端
  • YOLOv5m：平衡型（21.2M参数）
  • YOLOv5l/x：高精度型（46.5M/86.7M参数）

通过weights参数指定模型路径或预训练权重：

  1. model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)  # 方法1
  2. # 或
  3. from yolov5.models.experimental import attempt_load
  4. model = attempt_load('yolov5s.pt', map_location='cuda')  # 方法2

2. 自定义模型训练（进阶）

若需训练自定义数据集：

  1. 准备标注数据（YOLO格式：class x_center y_center width height）
  2. 修改data/coco.yaml配置文件
  3. 执行训练命令：
    1. python train.py --img 640 --batch 16 --epochs 50 --data coco.yaml --weights yolov5s.pt

四、推理流程实现

1. 输入数据预处理

支持多种输入类型：

  • 图像：PIL.Image或numpy.ndarray
  • 视频：cv2.VideoCapture
  • 实时流：RTSP/USB摄像头

预处理关键步骤：

  1. import cv2
  2. import numpy as np
  4. def preprocess(img_path, img_size=640):
  5.     # 读取图像
  6.     img = cv2.imread(img_path)
  7.     # 调整大小（保持宽高比）
  8.     h, w = img.shape[:2]
  9.     r = img_size / max(h, w)
  10.     new_h, new_w = int(h * r), int(w * r)
  11.     img = cv2.resize(img, (new_w, new_h))
  12.     # 填充至正方形
  13.     padded_img = np.ones((img_size, img_size, 3), dtype=np.uint8) * 114
  14.     padded_img[:new_h, :new_w] = img
  15.     # 转换为Tensor并归一化
  16.     img_tensor = torch.from_numpy(padded_img.transpose(2, 0, 1)).float() / 255.0
  17.     img_tensor = img_tensor.unsqueeze(0).to('cuda')  # 添加batch维度
  18.     return img_tensor, (h, w)

2. 推理执行

核心推理代码：

  1. def run_inference(model, img_tensor):
  2.     with torch.no_grad():
  3.         results = model(img_tensor)
  4.     return results
  6. # 示例调用
  7. img_tensor, (orig_h, orig_w) = preprocess('bus.jpg')
  8. results = run_inference(model, img_tensor)

3. 后处理与结果解析

YOLOv5返回results对象包含：

  • results.xyxy[0]：边界框坐标（x1,y1,x2,y2）
  • results.pred[0]：类别ID和置信度
  • results.pandas().xyxy[0]：DataFrame格式结果

可视化示例：

  1. from yolov5.utils.plots import Annotator
  3. def visualize(img, results, orig_hw):
  4.     h, w = orig_hw
  5.     annotator = Annotator(img, line_width=3, example=str(model.names))
  6.     for *box, conf, cls in results.xyxy[0]:
  7.         label = f'{model.names[int(cls)]} {conf:.2f}'
  8.         annotator.box_label(box, label, color=(255,0,0))
  9.     return annotator.img
  11. # 恢复原始尺寸
  12. resized_img = cv2.resize(img_tensor[0].cpu().numpy().transpose(1,2,0), (w, h))
  13. visualized_img = visualize(resized_img, results, (h, w))
  14. cv2.imwrite('output.jpg', visualized_img)

五、性能优化技巧

1. 硬件加速

  • GPU利用：确保torch.cuda.is_available()为True
  • TensorRT加速（NVIDIA GPU）：
    1. from yolov5.export import attempt_export
    2. attempt_export(model, 'yolov5s.trt')  # 生成TensorRT引擎

2. 批量推理

  1. batch_size = 4
  2. batch_imgs = [preprocess(f'img_{i}.jpg')[0] for i in range(batch_size)]
  3. batch_tensor = torch.cat(batch_imgs, dim=0)
  4. batch_results = model(batch_tensor)

3. 模型量化

  1. quantized_model = torch.quantization.quantize_dynamic(
  2.     model, {torch.nn.Linear}, dtype=torch.qint8
  3. )

六、典型应用场景

1. 实时摄像头检测

  1. cap = cv2.VideoCapture(0)  # 或RTSP流地址
  2. while True:
  3.     ret, frame = cap.read()
  4.     if not ret: break
  5.     # 预处理
  6.     img_tensor, _ = preprocess(frame)
  7.     # 推理
  8.     results = run_inference(model, img_tensor)
  9.     # 可视化
  10.     vis_frame = visualize(frame, results, frame.shape[:2])
  11.     cv2.imshow('Detection', vis_frame)
  12.     if cv2.waitKey(1) == 27: break  # ESC退出

2. 视频文件处理

  1. video_path = 'test.mp4'
  2. output_path = 'output.mp4'
  3. cap = cv2.VideoCapture(video_path)
  4. fps = cap.get(cv2.CAP_PROP_FPS)
  5. w, h = int(cap.get(3)), int(cap.get(4))
  6. fourcc = cv2.VideoWriter_fourcc(*'mp4v')
  7. out = cv2.VideoWriter(output_path, fourcc, fps, (w, h))
  9. while cap.isOpened():
  10.     ret, frame = cap.read()
  11.     if not ret: break
  12.     # 推理流程同上...
  13.     out.write(vis_frame)
  15. cap.release()
  16. out.release()

七、常见问题解决方案

  1. CUDA内存不足：

    • 减小img_size参数
    • 使用torch.cuda.empty_cache()清理缓存

  2. 模型加载失败：

    • 检查权重文件完整性（md5sum yolov5s.pt）
    • 确保PyTorch与CUDA版本匹配

  3. 检测精度低：

    • 调整conf_thres（置信度阈值，默认0.25）
    • 使用更大模型（如yolov5l.pt）

八、扩展应用建议

  1. 多任务学习：结合分类头实现检测+分类
  2. 嵌入式部署：使用ONNX Runtime或TFLite转换模型
  3. 分布式推理：通过torch.nn.DataParallel实现多卡并行

九、总结与展望

本文系统阐述了YOLOv5与PyTorch在Python中的物体检测推理实现，从环境配置到性能优化提供了完整解决方案。实际开发中，建议根据具体场景选择模型规模，并通过量化、TensorRT等技术进一步优化推理速度。随着PyTorch 2.0的发布，动态形状推理等新特性将为物体检测应用带来更多可能性。

附录：完整代码示例见GitHub仓库（示例链接），包含Jupyter Notebook实现和Docker部署方案。

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数
活动
咨询