使用 Docker 容器来安装和使用特定版本的 CUDA

引言

最近配置一个模型运行环境时,由于模型需要的一些依赖库包括tensorflow,CUDA的的等版本与系统原有cuda冲突,使用conda配置了半天还是没能解决,因此改用docker来配置环境,以达到与主机系统完全隔离的目的.

步骤

我这里需要的时tensorflow=2.16.1,所以以这个为例

1. 安装 Docker

安装如果提示启动失败则重启计算机

2. 安装 NVIDIA Docker

nvidia-docker 是一个允许 Docker 使用 NVIDIA GPU 的工具，确保它正确安装非常关键。

1)添加 NVIDIA Docker 的官方仓库，并安装 nvidia-docker2 包：

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2

2)重启 Docker 服务

sudo systemctl restart docker

3)验证 NVIDIA Docker 是否安装成功
通过运行带有 GPU 支持的简单测试来验证 nvidia-docker 是否正常工作：

docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

3. 查找兼容版本

访问 TensorFlow 官方 GitHub 或 TensorFlow 官网，查找 TensorFlow 2.16.1 版本的兼容 CUDA 和 cuDNN 版本。
TensorFlow 2.16.1 兼容 CUDA 11.x 和 cuDNN 8.x

4. 拉取Docker 镜像

TensorFlow 也提供了预装特定 CUDA 和 cuDNN 版本的官方 TensorFlow GPU 镜像。这些镜像已经设置好了所有必需的配置，可以直接使用。使用以下命令拉取带有特定 CUDA 和 cuDNN 的 TensorFlow 镜像：

docker pull tensorflow/tensorflow:2.16.1-gpu

5. 运行 Docker 容器

使用以下命令运行容器，并验证 CUDA 和 TensorFlow 的安装：

docker run --gpus all -it  tensorflow/tensorflow:2.16.1-gpu python

然后在 Python 提示符中尝试导入 TensorFlow 并检查 GPU：

import tensorflow as tf
print("GPU Available: ", tf.test.is_gpu_available())

运行完显示gpu数量则说明环境没有问题

6. 注意事项

1. 宿主机上需要先安装Docker。
2. 有时，官方 TensorFlow Docker 镜像可能不会精确标记每个 TensorFlow 的小版本（如 2.16.1），而是提供主版本（如 2.16）。