在支持 GPU 的 Docker 容器中运行 AI 模型

• 下载项目文件 - 3.8 MB

引言

容器技术，例如 Docker，极大地简化了软件的依赖管理和可移植性。在本系列文章中，我们将探讨 Docker 在机器学习 (ML) 场景中的应用。

本系列文章假设您熟悉 ML、容器化以及特别是 Docker。欢迎您下载 项目代码。

在 上一篇文章 中，我们为实验、训练和推理创建了基本的容器。为此，我们使用了 Intel/AMD 和 ARM CPU。在本文中，我们将利用 Nvidia GPU 的强大功能来减少训练和推理时间。

准备 Docker 主机以使用 Nvidia GPU

要从 Docker 使用 GPU，我们需要一个配备 Nvidia GPU 和 Linux 的主机（自 2020 年 12 月以来，GPU 支持也通过 适用于 Linux 的 Windows 子系统 (WSL2) 在 Windows 上运行）。在云端，您只需选择合适的 VM 大小和操作系统镜像。例如，Azure 上的 NC6 和带有 Ubuntu 18.04 的数据科学虚拟机。

根据 Linux 发行版和 GPU 型号的不同，本地机器上的配置 可能要求更高

1. 确保您的主机上已安装 Nvidia GPU 驱动程序。
2. 安装适合您的发行版的 nvidia-container-runtime。
3. 重启 Docker 守护进程。

现在您应该能够使用 --gpus 属性运行一个容器；例如，仅使用第一个 GPU

$ docker run --gpus "device=0" nvidia/cuda:11.2.1-runtime nvidia-smi

重要的是在主机和容器中使用相同的 CUDA 版本（在我们的例子中是 11.2）。如果版本不匹配，容器将无法启动，并显示类似“unsatisfied condition: cuda>=11.0.”的错误。

预测 Dockerfile

我们建议您始终从给定任务可用的最小基本镜像开始，通常后缀为“runtime”。虽然“devel”后缀似乎更合适，但它所表示的镜像包含许多在大多数 ML 场景中不需要的工具。

使用 Nvidia 提供的基本镜像（例如，上面提到的 nvidia/cuda:11.2.1-runtime），在此处安装 Python 和我们的库，然后完成任务，这可能很诱人。不幸的是，这行不通，至少对于 TensorFlow 而言。我们可以按照提供的 分步说明 或使用官方推荐的 Tensorflow Docker 镜像。

我们将采用后一种选择。除了 FROM 语句之外，新的预测 Dockerfile 的其余部分与我们用于仅 CPU 版本的相同

FROM tensorflow/tensorflow:2.3.2-gpu 
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update \
&& apt-get -y install --no-install-recommends ffmpeg libsm6 libxext6 \
&& apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/*
ARG USERNAME=mluser
ARG USERID=1000
RUN useradd --system --create-home --shell /bin/bash --uid $USERID $USERNAME
COPY requirements.txt /tmp/requirements.txt
RUN pip3 install --no-cache-dir -r /tmp/requirements.txt \
&& rm /tmp/requirements.txt
USER $USERNAME
WORKDIR /home/$USERNAME/app
COPY app /home/$USERNAME/app
ENTRYPOINT ["python", "predict.py"]

请注意，我们使用的是 TensorFlow 版本 2.3.2，而不是之前的 2.4.1。这主要是因为版本 2.4.1 的官方 Docker 镜像要大得多（>5.5GB，而所选镜像约为 3GB）。稍微旧的版本足以满足我们的需求。

因为我们使用包含 TensorFlow（以及匹配的 NumPy 版本）的基本镜像，所以我们的 requirements.txt 缩小为一行

opencv-python==4.5.1.48

下载 项目代码 后，我们可以构建镜像

$ docker build --build-arg USERID=$(id -u) -t mld05_gpu_predict .

如前所述，如果不需要 --build-arg USERID 参数，我们可以跳过它（尤其是在 Windows 上）。

训练 Dockerfile

因为我们希望不仅将 GPU 用于预测，还用于训练，所以我们需要引入一个额外的镜像定义 – Dockerfile.train

FROM mld05_gpu_predict:latest
ENTRYPOINT ["python", "train.py"]

我们只需将我们的预测镜像作为基础，并添加一个覆盖 ENTRYPOINT 为 train.py 脚本的单层。这是一个小的权衡，以避免增加镜像数量和代码重复。我们不介意这里的“latest”标签，因为我们完全控制着我们使用的基本镜像。

现在让我们构建它

$ docker build -t mld05_gpu_train -f 'Dockerfile.train' .

运行训练

让我们尝试使用 GPU 和 CPU 进行训练 - 以比较性能。

使用 GPU 训练

$ docker run -v $(pwd)/data:/home/mluser/data -v $(pwd)/models:/home/mluser/models \
 --rm --user $(id -u):$(id -g) --gpus "device=0" \
 mld05_gpu_train --model_path ../models/mnist_model.h5 --epochs 5

请注意传递给训练脚本的 --model_path 和 --epochs 参数。

要仅使用 CPU 训练，我们只需删除 --gpus "device=0" 参数

$ docker run -v $(pwd)/data:/home/mluser/data -v $(pwd)/models:/home/mluser/models \
 --rm --user $(id -u):$(id -g) \
 mld05_gpu_train --model_path ../models/mnist_model.h5 --epochs 5

如您所见，GPU 使训练速度翻了一番（从 CPU 上的每 epoch 18 秒到 GPU 上的每 epoch 8 秒）。不多，但我们在这里训练一个非常简单的模型。对于实际任务，您可以预期 5-10 倍的改进。

如果您的结果不同，请查看 TensorFlow 日志。如果它未能使用 GPU，您将看到类似“Could not load dynamic library (…).”的错误。

运行预测

在训练完我们的模型后，我们可以检查它在预测中的表现

$ docker run -v $(pwd)/data:/home/mluser/data -v $(pwd)/models:/home/mluser/models \
 --rm --user $(id -u):$(id -g) --gpus "device=0" \
 mld05_gpu_predict --images_path /home/mluser/data/test_mnist_images/*.jpg

摘要

在本文中，我们已经成功地使用带有 GPU 支持的 Docker 容器运行了训练和推理。我们准备好使用 Docker 解决基本的 ML 任务。在下一系列中，我们将扩展我们的知识，以解决在实际场景中使用 Docker 时更复杂的情况。