在本系列文章中，将探讨在机器学习（ML）场景中使用Docker的情况。本系列假设已经熟悉ML、容器化以及Docker。欢迎下载项目代码。
在前一篇文章中，为实验、训练和推理创建了基本的容器，使用了Intel/AMD和ARM CPU。在本文中，将利用Nvidia GPU的力量来减少训练和推理的时间。
要使用Docker中的GPU，需要一个带有Nvidia GPU的主机和Linux系统（自2020年12月以来，GPU支持也适用于Windows，通过Windows子系统Linux（WSL2））。在云中，只需要选择适当的VM大小和操作系统镜像。例如，在Azure上使用NC6和带有Ubuntu 18.04的数据科学虚拟机。
根据Linux发行版和GPU型号，本地机器的配置可能会更加复杂：
现在应该能够使用--gpus属性运行容器；例如，仅使用第一个GPU：
$ docker run --gpus "device=0" nvidia/cuda:11.2.1-runtime nvidia-smi
重要的是在主机和容器中使用相同的CUDA版本（在例子中是11.2）。如果版本不匹配，容器将无法启动，并出现错误，例如“unsatisfied condition: cuda>=11.0”。
建议始终从给定任务可用的最小基础镜像开始，通常带有“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上）。
因为希望不仅用于预测，还用于训练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上的18秒每周期到GPU上的8秒每周期）。虽然不多，但在这里训练的是一个非常简单的模型。在实际任务中，可以预期5-10倍的改进。
如果结果不同，请查看TensorFlow日志。如果它未能使用GPU，将看到错误，例如“Could not load dynamic library (…)”。