Deep Learning is very popular field to achive ML and AI tasks and there are a lot of frameworks for this and they are updated day by day by community. So, with that there are major releases on frame works after a short duration of time. Most of the deep learning frameworks user GPU for computation and they require Nvidia CUDA and CUDNN compatibility when working with Nvidia GPUs. So in this tutorial, we will discuss that how we can use multiple frameworks and their different versions using docker.

Docker helps to create OS-Level virtualization to run multiple code snippets and apps in different os and with different configurations. There are also different Deep Learning frameworks images are available which can be used with docker with just single command and you can code using that environment. To use GPU with docker, we need to setup nvidia-docker along docker. You can follow these urls to easily install docker and setup nvidia docker on machine.

Docker Setup (Ubuntu 20.04)

Nvidia Docker Setup

Once docker and nvidia docker is setup, you can test using this command is gpu is accessible. It pulls a nvidia cuda base repository and runs nvidia-smi command to show gpu details.

sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

Now, we can test different deep learning frameworks like Tensorflow, PyTorch etc.


Tensorflow is widely used for a lot of deep learning algorithms including computer vision, natural language processing, forecasting and many other language related works. There are different versions of tensorflow in development, beta and stable releases. So, instead of installing every version, we can use their docker images which has already tensorflow and other related packages setup and we can just pull and use it. For all latest releases and docker images check this docker hub url.

To pull and run a docker images, we can just pull using docker commands. Tensorflow also offers docker images with jupyter notebook already setup which helps to easily modify and run code inside container. So, lets pull a docker image with jupyter and run on system.

# if images is not on system, it will pull from docker hub.
# map port 8888 to 8888 inside container as jupter running on this port inside container.

docker run -p 8888:8888 --gpus all tensorflow/tensorflow:2.7.1-gpu-jupyter

Now open http://localhost:8888 in your browser or click on url directly with token to open jupyter notebook. Now you can use notebook to write and run code with specified tensorflow version. You can also install new packages with pip within jupyter cells using ! keyword like this !pip install opencv-python. Now docker container is running in isolated environment so we will need to use jupyter notebook upload and download feature to upload data to container instance and download results from it.

We can also mount volumes to docker container so that we specify a drive or path on our local system which can be used on docker container and on local system as well.

# volume
# local_path:container_path
docker run -p 8888:8888 -v /home/data:/tf/data --gpus all tensorflow/tensorflow:2.7.1-gpu-jupyter

This way we dont need to upload/download and avoid duplication of data and saves time.


Like tensorflow, Pytorch is widely used and is very popular for deep learning models development. It also offers docker container which can be utilized for development and deployment of models easily. Here is the minimal example to use pytorch gpu with docker.

docker run -p 8888:8888 --gpus all pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime

Pytorch does not offer any image with jupyter already setup. So, we need to either build image using current image or we can install jupyter in Pseudoterminal and we can run commands to install jupyter and then launch jupyter notebook.

docker run -p 8888:8888 --gpus all -it pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime

Now, we can use both linux and pip commands to install packages like jupyter.

pip install jupyter

# run notebook after installation
jupyter notebook

Build Docker Image

We can also use this pytorch image as base image and build a new docker image with new packages and set jupyter notebook to run on container startup. First create a dockerfile inside a directory and write this code to create a new docker image.

# Use pytorch with cuda image as base image
FROM pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime

# install python packages
RUN pip install jupyter

# specify port to run notebook

# start jupyter notebook
# NotebookApp.token is token for authentication (leave if you want to auto generate)
CMD exec jupyter notebook --ip --port $PORT --allow-root --NotebookApp.token='qw1234567'

Now, we build an image using docker build command and create a new image.

docker build -t pytorch-gpu-jupyter .

Once image is created, we can start a docker container using this command which will start a jupyter instance automatically.

docker run -p 8888:8888 -v /home/data:/data --gpus all pytorch-gpu-jupyter

Now, open http://localhost:8888 in browser and enter Notebook token defined in dockerfile and you can access notebook.

Import any packages installed or install new packages and can test gpu availablity using tensorflow or pytorch built in functions. For more information related to docker, view official docker documentation.