如何在Docker中安装Ultralytics？

Question

**问题：如何在Docker中安装Ultralytics并配置GPU支持？** 在使用Docker部署Ultralytics（如YOLOv8）时，常见的问题是：如何正确构建包含PyTorch和CUDA支持的镜像，并确保Ultralytics库能正常调用GPU进行训练或推理？用户常遇到的难点包括：基础镜像的选择、CUDA与CUDNN版本的匹配、依赖库的安装以及容器内GPU设备的访问权限配置。此外，如何通过Docker Compose简化部署流程也是一大关注点。本文将详细介绍如何编写Dockerfile、安装Ultralytics、配置GPU支持并验证安装是否成功。

Answer 1

1. 引言：为什么要在Docker中部署Ultralytics并启用GPU支持？

随着深度学习在计算机视觉领域的广泛应用，YOLO系列模型（尤其是YOLOv8）因其高效性和准确性受到广泛关注。Ultralytics作为YOLOv8的官方实现库，提供了完整的训练与推理接口。然而，在生产环境中，为了实现环境隔离和部署便捷性，通常会使用Docker进行容器化部署。

在容器中运行深度学习模型时，GPU支持至关重要。本文将从基础镜像选择、CUDA环境配置、Ultralytics安装到最终验证的完整流程进行详细说明，并提供Docker Compose的简化部署方案。

2. 环境准备与基础镜像选择

在构建Docker镜像之前，需确保宿主机已安装以下组件：

• Docker Engine
• NVIDIA Container Toolkit（用于GPU支持）
• CUDA驱动（版本需与后续选择的镜像兼容）

推荐使用NVIDIA提供的PyTorch官方镜像作为基础镜像，例如：

FROM pytorch/pytorch:2.0.1-cuda11-8-focal

该镜像已预装CUDA 11.8和CUDNN，适配PyTorch 2.0.1。选择镜像时应确保其CUDA版本与宿主机驱动版本兼容，避免因版本不匹配导致运行失败。

3. 编写Dockerfile并安装Ultralytics

下面是一个完整的Dockerfile示例，用于构建包含Ultralytics和GPU支持的镜像：

FROM pytorch/pytorch:2.0.1-cuda11-8-focal

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    ffmpeg \
    libgl1 \
    && rm -rf /var/lib/apt/lists/*

# 安装Ultralytics
RUN pip install --no-cache-dir ultralytics

# 设置工作目录
WORKDIR /workspace

# 挂载模型与数据目录
VOLUME ["/workspace/data", "/workspace/models"]

# 默认启动命令
CMD ["bash"]
    

该Dockerfile基于PyTorch官方镜像，安装了必要的系统依赖和Ultralytics库，并设置了工作目录与数据卷，便于后续训练或推理。

4. 配置容器GPU访问权限

为使容器能访问GPU资源，需在运行容器时指定--gpus参数：

docker run --gpus all -it --rm -v $(pwd):/workspace ultralytics-gpu

其中，--gpus all表示允许容器使用所有GPU资源，-v $(pwd):/workspace将当前目录挂载到容器的工作目录中。

此外，还需确保NVIDIA Container Toolkit已正确安装并配置。可通过以下命令验证：

nvidia-smi

若显示GPU信息，则说明环境配置成功。

5. 使用Docker Compose简化部署流程

为简化部署与管理，可使用Docker Compose编排服务。以下是一个docker-compose.yml文件示例：

version: '3.8'

services:
  yolov8:
    image: ultralytics-gpu
    build: .
    volumes:
      - .:/workspace
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - NVIDIA_DRIVER_CAPABILITIES=compute,utility
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    stdin_open: true
    tty: true
    

该配置文件定义了一个名为yolov8的服务，启用了GPU资源并挂载了本地目录，便于快速启动和调试。

6. 验证Ultralytics是否成功调用GPU

进入容器后，可运行以下Python脚本验证GPU是否可用：

import torch
from ultralytics import YOLO

# 检查CUDA是否可用
print("CUDA available:", torch.cuda.is_available())

# 加载YOLOv8模型
model = YOLO('yolov8s.pt')

# 推理测试（假设存在test.jpg图像）
results = model('test.jpg')
results.show()
    

若输出显示CUDA可用，并成功运行推理任务，则说明Ultralytics已正确配置GPU支持。