您可通过 Docker 快速在 GPU 实例上运行 TensorFlow,该方式仅需实例已安装 NVIDIA驱动程序,无需安装 NVIDIA CUDA 工具包。
前提条件
- GPU云主机已安装GPU 驱动。
- 本文操作步骤以CentOS操作系统为例。
操作步骤
1.安装 Docker
a. 在安装Docker新版本之前,请卸载任何此类旧版本以及关联的依赖项。
sudo yum remove docker \
docker-client \
docker-client-latest \
docker-common \
docker-latest \
docker-latest-logrotate \
docker-logrotate \
docker-engine
b. 设置Docker 存储库。
sudo yum install -y yum-utils
sudo yum-config-manager --add-repo https://download.docker.com/linux/centos/docker-ce.repo
c. 安装Docker 引擎。
sudo yum install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
d. 启动docker。
sudo systemctl start docker
2.安装 TensorFlow
a. 设置 NVIDIA 容器工具包
- 设置存储库和GPG密钥。
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo
- 更新包列表后安装nvidia-container-toolkit包(和依赖项)。
sudo yum clean expire-cache
sudo yum install -y nvidia-container-toolkit
- 配置Docker守护程序以识别 NVIDIA 容器运行时。
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
- 通过运行基本CUDA容器来测试工作设置。
sudo docker run --rm --runtime=nvidia --gpus all nvidia/cuda:11.6.2-base-ubuntu20.04 nvidia-smi
b.下载TensorFlow Docker镜像
官方 TensorFlow Docker 镜像位于 tensorflow/tensorflow Docker Hub 代码库中。镜像版本按照以下格式进行 标记:
标记 | 说明 |
---|---|
latest | 标签包含最新版本(不包括候选版本、alpha 和 beta 等预发布版本) |
nightly | 附带最新的TensorFlow nightly Python 包 |
version | 指定TensorFlow 二进制镜像的版本,例如 2.13.0 |
devel | 标签不再支持,请改用TensorFlow SIG Build Dockerfiles |
custom-op | 标签不再支持,请改用TensorFlow SIG Build Dockerfiles |
每个基本标记都有会添加或更改功能的变体:
标记变体 | 说明 |
---|---|
tag-gpu | 标签基于Nvidia CUDA。您需要nvidia-docker来运行它们 |
tag-jupyter | 标签包括Jupyter 和一些 TensorFlow 教程笔记本 |
您可以一次使用多个变体。例如,以下命令会将 TensorFlow 版本镜像下载到计算机上。
docker pull tensorflow/tensorflow # latest stable release
docker pull tensorflow/tensorflow:nightly-gpu # nightly release w/ GPU support
docker pull tensorflow/tensorflow:latest-gpu-jupyter # latest release w/ GPU support and Jupyter
c.启动TensorFlow Docker容器
启动配置 TensorFlow 的容器,请使用以下命令格式。如需了解更多信息,请参见 Docker run reference。
docker run [-it] [--rm] [-p hostPort:containerPort] tensorflow/tensorflow[:tag] [command]
示例
1.使用仅支持 CPU 的镜像的示例
如下所示,使用带latest 标记的镜像验证 TensorFlow 安装效果。Docker 会在首次运行时下载新的 TensorFlow 镜像:
docker run -it --rm tensorflow/tensorflow python
其他 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=...。
2.使用支持 GPU 的镜像的示例
a. 执行以下命令,下载并运行支持 GPU 的 TensorFlow 镜像,并使用 Python 解释器启动。
docker run -it --rm --runtime=nvidia tensorflow/tensorflow:latest-gpu python
b. 设置支持 GPU 镜像可能需要一段时间。如果重复运行基于 GPU 的脚本,您可以使用 docker exec 重复使用容器。执行以下命令,使用最新的 TensorFlow GPU 镜像在容器中启动 bash shell 会话:
docker run --gpus all -it tensorflow/tensorflow:latest-gpu bash