CUDA TensorFlow and PyTorch Docker image

An Ubuntu 18.04 NVIDIA Docker image with CUDA 10.1 and CuDNN 7 with TensorFlow and PyTorch

Contains:

• TensorFlow v2.1.0
• PyTorch v1.4.0

Installation

• Make sure that you have the nvidia-container-toolkit installed on the host machine.
• Check the list of available tags on Docker Hub to find the tag you want.
• Use docker pull to pull down the image corresponding to the tag. For example:

docker pull matthewfeickert/cuda-tf-torch:cuda10.1-cudnn7-runtime-ubuntu18.04

Use

To check that NVIDIA GPUS are being properly detected run

docker run --rm --gpus all matthewfeickert/cuda-tf-torch:cuda10.1-cudnn7-runtime-ubuntu18.04 nvidia-smi

and check if the nvidia-smi output appears correctly.

To run (interactively) using GPUs on the host machine:

docker run --rm -it --gpus all matthewfeickert/cuda-tf-torch:cuda10.1-cudnn7-runtime-ubuntu18.04