使用 Docker 和 Ultralytics YOLOv8 训练自定义目标检测模型全流程

1 天 ago

木它

1 minute

前言

最近在学习目标检测，想用 Ultralytics 的 YOLOv8 框架训练自己拍摄的“蓝色黄色小球”数据集。由于本地环境复杂，决定用 Docker 来隔离环境，保证依赖一致性。本文记录整个流程，便于以后复用和查阅。

1. 准备工作

1.1 安装 Docker

请确保电脑已安装 Docker Desktop，并且已开启 GPU 支持（NVIDIA 驱动 & CUDA 已配置好）。

1.2 检查 GPU 驱动

在命令行（PowerShell）输入：

nvidia-smi

如果能看到显卡信息，说明驱动没问题。

2. 数据集准备

我的数据集结构如下：

E:\
└── 机器学习训练集\
    └── 蓝色黄色小球\
        └── boll\
            ├── train\
            │    ├── images\
            │    └── labels\
            └── val\
                 ├── images\
                 └── labels\

images 目录下是图片（jpg/png等）
labels 目录下是和图片同名的 txt 文件（YOLO 格式标注）

3. 使用 Docker 启动 YOLO 环境并挂载数据集

在 PowerShell 中运行如下命令，将本地数据集挂载到 Docker 容器内：

docker run -it --ipc=host --gpus all `
  -v "E:/机器学习训练集/蓝色黄色小球/boll:/ultralytics/data/boll" `
  ultralytics/ultralytics:latest

-v 参数用于挂载本地目录到容器内部。
挂载后，容器内的数据路径为 /ultralytics/data/boll。

4. 准备 data.yaml 配置文件

在本地 E:/机器学习训练集/蓝色黄色小球/boll/ 下新建 data.yaml 文件，内容如下：

train: /ultralytics/data/boll/train/images
val: /ultralytics/data/boll/val/images

nc: 2
names: ['blue_ball', 'yellow_ball']

nc：类别数
names：类别名称，顺序要和你的标注一致
train 和 val 路径要用容器内的挂载路径

5. 在容器里开始训练

进入容器后，执行以下命令开始训练：

yolo task=detect mode=train model=yolov8n.pt data=/ultralytics/data/boll/data.yaml epochs=100 imgsz=640

参数说明：

model=yolov8n.pt：使用 YOLOv8 nano 预训练模型
data：指向 data.yaml 文件
epochs：训练轮数
imgsz：输入图片尺寸

6. 查看训练结果

训练完成后，模型和日志会保存在 /ultralytics/runs/detect/ 目录下。

可以通过如下命令查看：

ls /ultralytics/runs/detect/

如果需要把结果拷贝回本地，可用 docker cp 命令。

7. 常见问题总结

路径中有中文或空格：建议尽量用英文路径，或确保 Docker Desktop 已授权对应盘符。
labels 格式：务必是 YOLO 标准格式，类别编号从 0 开始。
挂载后目录为空：请检查路径拼写和 Docker 权限设置。

8. 参考资料

9. 用命令行直接导出 YOLO 模型为 ONNX

best.pt 路径确保正确：

yolo export model=/ultralytics/data/boll/train3/weights/best.pt format=onnx

model=... 指定你的 .pt 文件路径
format=onnx 指定导出格式为 ONNX

导出后会在 runs/ 目录下生成对应的 ONNX 文件。

还可以加其他参数，比如指定输出路径、导出的输入尺寸等：

yolo export model=/ultralytics/data/boll/train3/weights/best.pt format=onnx imgsz=640

更多参数可以用 yolo export --help 查看。