在Docker中使用TensorFlow（GPU版+Ubuntu 20.04）

Posted 2022-11-26 白马负金羁

在此前的文章【1】中，我们已经介绍了如何在Ubuntu 20.04中安装Docker。本文将介绍如何在其中使用TensorFlow（GPU）版。

一、安装NVIDIA-Docker

安装完成Docker并检查安装正确（能跑出来hello-world）后，如果需要Docker容器中有GPU支持，需要再安装NVIDIA-Docker，同样参考该项目的官方文档【2】。

Docker-CE on Ubuntu can be setup using Docker’s official convenience script:

Setting up NVIDIA Container Toolkit——Setup the package repository and the GPG key:

 然后，Install the nvidia-docker2 package (and dependencies) after updating the package listing:

$ sudo apt-get update
$ sudo apt-get install -y nvidia-docker2

Restart the Docker daemon to complete the installation after setting the default runtime:

$ sudo systemctl restart docker

 At this point, a working setup can be tested by running a base CUDA container。用最新的CUDA镜像来测试一下nvidia-smi（检查是否安装成功，安装成功，则会显示关于GPU的信息）：

二、用Docker（NVIDIA-Docker）载入TensorFlow镜像

Start a CPU-only container，可以使用：

$ docker run -it --rm tensorflow/tensorflow bash

Start a GPU container, using the Python interpreter（其中，The "--gpus all" makes all the gpus available on the host to be visible to the container.）：

$ docker run -it --gpus all --rm --runtime=nvidia tensorflow/tensorflow:latest-gpu python

Run a Jupyter notebook server with your own notebook directory (assumed here to be ~/notebooks).

• 1.xx-, latest-, and nightly- tags come with TensorFlow pre-installed. Versioned tags contain their version, the latest- tags contain the latest release (excluding pre-releases like release candidates, alphas, and betas), and the nightly images come with the latest TensorFlow nightly Python package.
• -gpu tags are based on Nvidia CUDA.
• -jupyter tags include Jupyter and some TensorFlow tutorial notebooks.

注意上图中 To access the notebook, open ...

所以有几个方法，你可以打开一个Jupyter Notebook。

然后在其中执行

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

就会输入Available GPUs的数量。

这是确定你的TensorFlow是用了GPU的一种方法，参见【3】。

---

参考文献与推荐阅读材料

【1】在Ubuntu 20.04上安装Python虚拟环境

【2】Installation Guide：Docker

【3】Nvidia TensorFlow