Running code using Docker

Docker uses containers to create virtual environments that isolate a TensorFlow installation from the rest of the system. TensorFlow programs are run within this virtual environment that can share resources with its host machine (access directories, use the GPU, connect to the Internet, etc.). The TensorFlow Docker images are tested for each release.

Docker is the easiest way to enable TensorFlow GPU support on Linux since only the NVIDIA® GPU driver is required on the host machine (the NVIDIA® CUDA® Toolkit does not need to be installed).

TensorFlow Docker requirements

1. Install Docker on your local host machine.
2. For GPU support on Linux, install NVIDIA Docker support.
   • Take note of your Docker version with docker -v. Versions earlier than 19.03 require nvidia-docker2 and the --runtime=nvidia flag. On versions including and after 19.03, you will use the nvidia-container-toolkit package and the --gpus all flag. Both options are documented on the page linked above.

Note: To run the docker command without sudo, create the docker group and add your user. For details, see the post-installation steps for Linux.

Download a TensorFlow Docker image

The official TensorFlow Docker images are located in the tensorflow/tensorflow Docker Hub repository. Image releases are tagged using the following format:

You can use multiple variants at once. For example, the following downloads TensorFlow release images to your machine:

Examples using CPU-only images Let's verify the TensorFlow installation using the latest tagged image. Docker downloads a new TensorFlow image the first time it is run:

GPU support

Docker is the easiest way to run TensorFlow on a GPU since the host machine only requires the NVIDIA® driver (the NVIDIA® CUDA® Toolkit is not required).

Install the Nvidia Container Toolkit to add NVIDIA® GPU support to Docker. nvidia-container-runtime is only available for Linux. See the nvidia-container-runtime platform support FAQ for details.

Check if a GPU is available:

lspci | grep -i nvidia

Verify your nvidia-docker installation:

docker run --gpus all --rm nvidia/cuda nvidia-smi

Note: nvidia-docker v2 uses --runtime=nvidia instead of --gpus all. nvidia-docker v1 uses the nvidia-docker alias, rather than the --runtime=nvidia or --gpus all command line flags. Examples using GPU-enabled images

Download and run a GPU-enabled TensorFlow image (may take a few minutes):

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])))"

It can take a while to set up the GPU-enabled image. If repeatably running GPU-based scripts, you can use docker exec to reuse a container.