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

目录

介绍

准备Docker主机以使用Nvidia GPU

预测Dockerfile

训练Dockerfile

运行训练

运行预测

总结

---

• 下载项目文件 - 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支持也可通过Windows Subsystem for Linux (WSL2)在Windows上运行）。在云中，您需要做的就是选择合适的VM大小和操作系统镜像。例如，NC6和数据科学虚拟机与Azure上的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）。如果版本不匹配，容器将无法启动，并出现“不满足条件：cuda>=11.0”之类的错误。

预测Dockerfile

我们建议您始终从可用于给定任务的最小基础镜像开始，通常后缀为“运行时”。虽然“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 .

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

训练Dockerfile

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

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

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

现在让我们构建它：

$ 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，您将看到类似“无法加载动态库 (...)”的错误消息。

运行预测

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

$ 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时很典型。

https://www.codeproject.com/Articles/5300729/Running-AI-Models-in-GPU-Enabled-Docker-Containers