分享

在启用GPU的Docker容器中运行AI模型

 码农9527 2021-05-06

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

    本系列假定您熟悉ML,一般的容器化,尤其是Docker。欢迎您下载项目代码。

    在上一篇文章中,我们创建了用于实验,训练和推断的基本容器。为此,我们使用了Intel/AMD和ARMCPU。在这一节中,我们将利用NvidiaGPU的功能来减少训练和推理时间。

    准备Docker主机以使用NvidiaGPU

    要使用来自Docker的GPU,我们需要具有NvidiaGPU和Linux的主机(自2020年12月以来,GPU支持也通过WindowsLinux子系统(WSL2)在Windows上运行)。在云中,您所需要做的就是选择适当的VM大小和OS映像。例如,在Azure上具有Ubuntu18.04的NC6和数据科学虚拟机。

    根据Linux发行版和GPU模型的不同,在本地计算机上的配置可能会要求更高:

    确保主机上已安装NvidiaGPU驱动程序。

    安装适合您的发行版的nvidia-container-runtime。

    重新启动Docker守护程序。

    现在,您应该可以运行带有--gpus属性的容器了。例如,仅使用第一个GPU:

$ docker run --gpus "device=0" nvidia/cuda:11.2.1-runtime nvidia-smi1复制代码类型:[html]

WEB

    在主机和容器中使用相同的CUDA版本(在本例中为11.2)很重要。如果版本不匹配,则容器将无法启动,并显示诸如“条件不满足:cuda>=11.0”之类的错误。

    预测Dockerfile

    我们建议您始终从给定任务可用的最小基础映像开始,通常后缀为“运行时”。尽管“devel”后缀似乎更合适,但它表示的图像包含了大多数ML场景中不需要的许多工具。

    使用Nvidia提供的基本映像(例如,上面提到的nvidia/cuda:11.2.1-runtime),在其中安装Python和我们的库可能很诱人。不幸的是,这至少在TensorFlow中不起作用。我们可以按照提供的逐步说明进行操作,也可以使用官方推荐的TensorflowDocker映像。

    我们将使用后一种选择。除了该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"]123456789101112131415复制代码类型:[html]

    请注意,我们使用TensorFlow版本2.3.2,而不是以前的2.4.1。主要原因是2.4.1版的官方Docker映像要大得多(>5.5GB,而选定的映像约为3GB)。稍旧的版本足以满足我们的目的。

    因为我们使用已包含TensorFlow的基础图像(以及匹配的NumPy版本),所以我们的requirements.txt缩小为一行:

opencv-python==4.5.1.481复制代码类型:[html]

    下载项目代码后,我们可以构建图像:

$ docker build --build-arg USERID=$(id -u) -t mld05_gpu_predict .1复制代码类型:[html]

    如前所述,--build-argUSERID如果不需要该参数,则可以跳过该参数(尤其是在Windows上)。

    培训Dockerfile

    因为我们不仅希望将GPU用于预测,还希望将其用于训练,所以我们需要引入其他图像定义-Dockerfile.train:

FROM mld05_gpu_predict:latest
ENTRYPOINT ["python", "train.py"]12复制代码类型:[html]

    我们ENTRYPOINT仅以预测图像为基础,并添加一个用train.py脚本覆盖的单层。为了避免增加图像和代码重复的数量,这是一个很小的折衷。我们不介意此处的“最新”标签,因为我们可以完全控制所使用的基本图像。

    现在,我们来构建它:

$ docker build -t mld05_gpu_train -f 'Dockerfile.train' .1复制代码类型:[html]

    进行培训

    让我们尝试使用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 5123复制代码类型:[html]

WEB

    请注意--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 5123复制代码类型:[html]

WEB

    如您所见,GPU将训练速度提高了一倍(从CPU上的每个周期18秒增加到GPU上的每个周期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/*.jpg123复制代码类型:[html]

WEB

    概括

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

    转藏 分享 献花(0

    0条评论

    发表

    请遵守用户 评论公约

    类似文章 更多