YOLOv5 Docker 镜像操作指南

这是一个基于 YOLOv5 v7.0 构建的 YOLO Docker 镜像，适用于 YOLOv5 v7.0版本的 AI 视觉处理任务。通过该镜像，用户可以轻松进行 图像分类、目标检测、目标跟踪、姿态识别、图像分割 等任务。该镜像支持以下操作：

• 模型训练（train）：在自定义数据集上训练 YOLO模型。
• 模型验证（val）：验证模型性能，获得关键指标。
• 模型推理（inf）：对图像、视频流进行目标检测和分类。
• 部署优化（opt）：通过 TensorRT 等工具，优化 GPU 加速推理。

1. 环境配置

• 核心硬件：NVIDIA GeForce RTX 4090
• 操作系统：22.04.4 LTS (Jammy Jellyfish)
• CUDA 版本：12.4
• cuDNN 版本：90100
• Python 版本：3.10.12
• Torch 版本：2.3.0a0+6ddf5cf85e.nv24.04
• TorchVision 版本：0.18.0a0
• TensorRT 版本：8.6.3
• YOLOv5 版本：v7.0

2. 数据集准备与模型训练

准备自定义数据集

• 支持的数据集格式：COCO、VOC 等，也可以使用自定义数据集。
• 使用标注工具 labelImg 进行数据标注。
• 镜像已包含提供 COCO 数据集。

模型训练

• 使用以下命令进行模型训练：

  python train.py --img 640 --batch 16 --epochs 100 --data coco.yaml --weights yolov5s.pt
  

• 可根据需求调整超参数（如学习率、批次大小）以优化训练效果。

模型验证与评估

• 使用验证集评估模型性能：

  python val.py --weights yolov5s.pt --data coco.yaml --img 640
  

3. 模型优化

• 环境中预装了 YOLOv5 的多种格式的模型：PyTorch、ONNX 和 TensorRT。

• 如果您训练了自定义模型，可将其转换为所需格式，如 TensorRT 格式。参考以下命令将 PyTorch 格式的模型导出为 ONNX 或 TensorRT 格式：

  python export.py --weights yolov5s.pt --include onnx
  

  或

  # TensorRT 格式
  trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine
  

4. 实际应用示例

静态图片分类推理

• 使用以下命令对静态图片进行分类推理：

  python detect.py --weights yolov5s.pt --source path/to/image.jpg
  

静态图片目标检测

• 使用以下命令检测静态图片中的目标：

  python detect.py --weights yolov5s.pt --source path/to/image.jpg
  

视频流目标检测

• 对实时视频流进行检测：

  python detect.py --weights yolov5s.pt --source path/to/video.mp4
  

5. 常见问题与故障排除

模型加载错误

• 确保模型与 PyTorch 版本兼容，并检查 CUDA 驱动是否正确安装。

推理速度较慢

• 使用较轻量化的模型（如 YOLOv5s）以提高推理速度。

容器网络问题

• 确保容器的网络配置正确，使用 docker network 工具进行排查。

模型格式问题

• 建议使用 YOLOv5 原生代码导出模型格式，而非依赖 TensorRT 工具或其他方法直接导出。

6. 参考资料

• YOLOv5 官方文档
• Docker 官方文档
• NVIDIA 容器工具链文档