YOLOv8 Docker 镜像操作指南

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

• 模型训练（train）：在自定义数据集上训练 YOLOv8 及以上版本模型。
• 模型验证（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
• Ultralytics版本：8.3.24

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

准备自定义数据集

• 支持的数据集格式：COCO、VOC 等，也可以使用自己的数据集。
• 使用标注工具 labelImg 进行数据标注。
• 镜像提供了coco数据集下载脚本，执行以下命令下载COCO数据集。

  sh download_coco_dataset.sh
  

模型训练

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

  yolo train model=yolov8n.pt data=coco.yaml epochs=100
  

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

模型验证与评估

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

  yolo val model=yolov8n.pt data=coco.yaml
  

3. 模型优化

• 环境已经预先准备好了YOLOv8的三种格式的模型：PyTorch、ONNX和TensorRT。

• 如果您已经训练好了自己的模型，可以通过模型导出的方式将其格式转换为自己需要的格式，如TensorRT格式。

• 参考使用以下命令将PyTorch格式模型导出为ONNX或TensorRT格式。

  yolo export model=yolov8n.pt format=onnx
  

  或

  # TensorRT格式
  yolo export model=yolov8n.pt format=tensorrt
  

  建议：使用Python代码来完成模型格式转换，这样可以避免模型转换错误。

  from ultralytics import YOLO
  
  # 设置模型的路径
  model_path = yolov8n.pt  # 替换为你的 YOLOv8 模型文件路径
  onnx_output_path = yolov8.onnx  # 输出的 ONNX 文件路径
  
  # 加载模型
  model = YOLO(model_path)
  
  # 导出模型为 ONNX 格式
  #model.export(format=onnx, simplify=True, opset=13)
  
  # 导出模型为 TensorRT 格式
  model.export(format=tensorrt, simplify=True, opset=13)
  
  

4. 实际应用示例

静态图片分类推理

• 运行以下命令对静态图片进行图片分类：

  yolo predict classify model=yolov8n-cls.pt source=path/to/image.jpg
  

静态图片目标检测

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

  yolo predict detect model=yolov8n.pt source=path/to/image.jpg
  

视频流目标检测

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

  yolo predict detect model=yolov8n.pt source=path/to/video.mp4
  

视频流图像分割

• 对实时视频流进行图像分割：

  yolo predict segment model=yolov8n-seg.pt source=path/to/video.mp4
  

视频流姿态识别

• 对实时视频流进行目标跟踪：

  yolo predict pose model=yolov8n-pose.pt source=path/to/video.mp4
  

视频流目标跟踪

• 对实时视频流进行目标跟踪：

  yolo track source=path/to/video.mp4
  

5. 常见问题与故障排除

模型加载错误

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

推理速度较慢

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

容器网络问题

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

模型格式问题

• 建议使用YOLO原生代码转换模型格式，而非TensorRT命令或yolo命令导出模型。

6. 参考资料

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