在Docker容器中使用NVIDIA GPU

在现代计算环境中,GPU加速对于执行深度学习和其他计算密集型任务至关重要。然而,当使用Docker容器时,GPU通常不会被自动配置。幸运的是,NVIDIA提供了一个名为NVIDIA Container Toolkit的工具,可以帮助解决这个问题。本文将详细介绍如何在Docker容器中访问和使用机器上的NVIDIA GPU。

安装NVIDIA GPU驱动

在开始之前,确保机器上已经安装了NVIDIA GPU驱动。这个过程可能因操作系统、NVIDIA GPU型号和驱动版本而异。安装完成后,可以通过运行nvidia-smi命令来验证驱动是否正确安装。如果命令成功执行并显示了GPU信息,那么机器已经成功识别了GPU。

使用NVIDIA Container Toolkit

NVIDIA Container Toolkit是一个特殊的Docker镜像,它能够自动检测机器上的GPU驱动,并将这些驱动传递给运行中的Docker容器。这样,就可以在容器内部运行nvidia-smi命令,并且所有的程序都能够访问GPU。

要使用NVIDIA Container Toolkit,需要在Dockerfile的顶部拉取相应的镜像,例如:

FROM nvidia/cuda:10.2-base CMD nvidia-smi

然后,使用以下命令构建镜像:

docker build . -t nvidia-test

最后,使用--gpus all标志运行容器:

docker run --gpus all nvidia-test

这样,Docker容器就能够看到GPU驱动,并且可以利用GPU进行计算了。

完整的Dockerfile示例

以下是一个完整的Dockerfile示例,它展示了如何使用NVIDIA Container Toolkit来构建一个深度学习模型训练和评估的容器

FROM nvidia/cuda:10.2-base CMD nvidia-smi # 设置环境 RUN apt-get update && apt-get install --no-install-recommends --no-install-suggests -y curl RUN apt-get install unzip RUN apt-get -y install python3 RUN apt-get -y install python3-pip COPY app/requirements_verbose.txt /app/requirements_verbose.txt RUN pip3 install -r /app/requirements_verbose.txt # 将应用从本地路径复制到容器路径 COPY app/ /app/ WORKDIR /app ENV NUM_EPOCHS=10 ENV MODEL_TYPE='EfficientDet' ENV DATASET_LINK='HIDDEN' ENV TRAIN_TIME_SEC=100 CMD ["python3", "train_and_eval.py"]

这个Docker容器将根据指定的参数训练和评估一个深度学习模型,同时利用基础机器的GPU

其他方法

除了使用NVIDIA Container Toolkit,还有其他方法可以将GPU驱动暴露给Docker,例如通过在构建镜像时包含配置GPU的命令。这种方法的缺点是每次重建镜像时都需要重新安装驱动,这会降低开发速度。此外,如果将Docker镜像迁移到具有不同GPU、操作系统或驱动的机器上,需要为每台机器重新编码这些步骤。

这种方法在Dockerfile中可能看起来像这样:

FROM ubuntu:14.04 MAINTAINER Regan <http://stackoverflow.com/questions/25185405/using-gpu-from-a-docker-container> RUN apt-get update && apt-get install -y build-essential RUN apt-get --purge remove -y nvidia* ADD ./Downloads/nvidia_installers /tmp/nvidia > 获取用来在主机上安装CUDA和NVIDIA驱动的安装文件 RUN /tmp/nvidia/NVIDIA-Linux-x86_64-331.62.run -s -N --no-kernel-module > 安装驱动。 RUN rm -rf /tmp/selfgz7 > 删除临时文件 RUN /tmp/nvidia/cuda-linux64-rel-6.0.37-18176142.run -noprompt > CUDA驱动安装程序。 RUN /tmp/nvidia/cuda-samples-linux-6.0.37-18176142.run -noprompt -cudaprefix=/usr/local/cuda-6.0 > CUDA样本 RUN export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64 > 添加CUDA库到PATH RUN touch /etc/ld.so.conf.d/cuda.conf > 更新ld.so.conf.d目录 RUN rm -rf /temp/* > 删除安装文件。

这种方法需要在本地文件夹中拥有NVIDIA驱动。可以将上述示例中的"./Downloads"文件夹替换为保存GPU驱动的目录。

现在知道了如何使用NVIDIA Container Toolkit将GPU驱动暴露给运行中的Docker容器。应该能够使用GPU驱动并运行Docker Compose而不会遇到任何问题。

沪ICP备2024098111号-1
上海秋旦网络科技中心:上海市奉贤区金大公路8218号1幢 联系电话:17898875485