使用 Docker 安装 TensorFlow 并设置 GPU/CPU 支持

GPU 云服务器（Cloud GPU Service）是基于 GPU 的快速、稳定、弹性的计算服务，主要应用于深度学习训练/推理、图形图像处理以及科学计算等场景。 GPU 云服务器提供和标准 CVM 云服务器一致的方便快捷的管理方式。GPU 云服务器通过其强大的快速处理海量数据的计算性能，有效解放用户的计算压力，提升业务处理效率与竞争力。

操作场景

您可通过 Docker 快速在 GPU 实例上运行 TensorFlow，且该方式仅需实例已安装 NVIDIA® 驱动程序，无需安装 NVIDIA® CUDA® 工具包。

本文介绍如何在 GPU 云服务器上，使用 Docker 安装 TensorFlow 并设置 GPU/CPU 支持。

说明事项

本文操作步骤以 Ubuntu 20.04 操作系统的 GPU 云服务器为例。

您的 GPU 云服务器实例需已安装 GPU 驱动。

说明

建议使用公共镜像创建 GPU 云服务器。若选择公共镜像，则勾选后台自动安装 GPU 驱动即可预装相应版本驱动。该方式仅支持部分 Linux 公共镜像，详情请参见 各实例支持的 GPU 驱动版本及安装方式。

操作步骤

安装 Docker

1. 登录实例，依次执行以下命令，安装所需系统工具。

sudo apt-get update

sudo apt-get install \

ca-certificates \

curl \

gnupg \

lsb-release

2. 执行以下命令，安装 GPG 证书，写入软件源信息。

sudo mkdir -p /etc/apt/keyrings

curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg

echo \

"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \

$(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

3. 依次执行以下命令，更新并安装 Docker-CE。

sudo apt-get update

sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin

安装 TensorFlow

设置 NVIDIA 容器工具包

1. 执行以下命令，设置包存储库和 GPG 密钥。详细信息请参见 Setting up NVIDIA Container Toolkit。

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \

&& curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \

&& curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \

sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \

sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

2. 执行以下命令，安装 nvidia-docker2 包及依赖项。

sudo apt-get update

sudo apt-get install -y nvidia-docker2

3. 执行以下命令，设置默认运行时重启 Docker 守护进程完成安装。

sudo systemctl restart docker

4. 此时可执行以下命令，通过运行基本 CUDA 容器来测试工作设置。

sudo docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi

返回结果如下所示：

+-----------------------------------------------------------------------------+

| NVIDIA-SMI 450.51.06 Driver Version: 450.51.06 CUDA Version: 11.0 |

|-------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|===============================+======================+======================|

| 0 Tesla T4 On | 00000000:00:1E.0 Off | 0 |

| N/A 34C P8 9W / 70W | 0MiB / 15109MiB | 0% Default |

| | | N/A |

+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=============================================================================|

| No running processes found |

+-----------------------------------------------------------------------------+

下载 TensorFlow Docker 镜像

官方 TensorFlow Docker 镜像位于 tensorflow/tensorflow Docker Hub 代码库中。镜像版本按照以下格式进行标记：

每个基本标记都有会添加或更改功能的变体：

您可以一次使用多个变体。例如，以下命令会将 TensorFlow 版本镜像下载到计算机上：

docker pull tensorflow/tensorflow # latest stable release

docker pull tensorflow/tensorflow:devel-gpu # nightly dev release w/ GPU support

docker pull tensorflow/tensorflow:latest-gpu-jupyter # latest release w/ GPU support and Jupyter

启动 TensorFlow Docker 容器

启动配置 TensorFlow 的容器，请使用以下命令格式。

docker run [-it] [--rm] [-p hostPort:containerPort] tensorflow/tensorflow[:tag] [command]

示例

使用仅支持 CPU 的镜像的示例

如下所示，使用带 latest 标记的镜像验证 TensorFlow 安装效果。Docker 会在首次运行时下载新的 TensorFlow 镜像：

docker run -it --rm tensorflow/tensorflow \

python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

其他 TensorFlow Docker 方案示例如下：

在配置 TensorFlow 的容器中启动 bash shell 会话：

docker run -it tensorflow/tensorflow bash

如需在容器内运行在主机上开发的 TensorFlow 程序，请通过 -v hostDir:containerDir -w workDir 参数，装载主机目录并更改容器的工作目录。示例如下：

docker run -it --rm -v $PWD:/tmp -w /tmp tensorflow/tensorflow python ./script.py

说明

向主机公开在容器中创建的文件时，可能会出现权限问题。通常情况下，最好修改主机系统上的文件。

使用 nightly 版 TensorFlow 启动 Jupyter 笔记本服务器：

docker run -it -p 8888:8888 tensorflow/tensorflow:nightly-jupyter

请参考 Jupyter 官网 相关说明，使用浏览器访问 http://127.0.0.1:8888/?token=...。

使用支持 GPU 的镜像的示例

执行以下命令，下载并运行支持 GPU 的 TensorFlow 镜像。

docker run --gpus all -it --rm tensorflow/tensorflow:latest-gpu \

python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

设置支持 GPU 镜像可能需要一段时间。如果重复运行基于 GPU 的脚本，您可以使用 docker exec 重复使用容器。 执行以下命令，使用最新的 TensorFlow GPU 镜像在容器中启动 bash shell 会话：

docker run --gpus all -it tensorflow/tensorflow:latest-gpu bash