This post describes how I set up a Docker container with Tensorflow/Keras, GPU support, and the Jupyter notebook, in 10 easy steps!
I’m running Kubuntu 20.04, which, for the purposes of this post, is the same as Ubuntu. The computer is an AMD ThreadRipper with an RTX-2070 video card that has GPUs.
Here’s how I got TensorFlow working.
sudo apt install nvidia-driver-455 nvidia-utils-455
curl https://get.docker.com | sh && sudo systemctl --now enable docker
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
&& curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list \
| sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2 && sudo systemctl restart docker
sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
You should see something like this:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06 Driver Version: 450.51.06 CUDA Version: 11.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 On | 00000000:00:1E.0 Off | 0 |
| N/A 34C P8 9W / 70W | 0MiB / 15109MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+
docker pull tensorflow/tensorflow:latest-gpu-jupyter
cd ~
mkdir -p docker/dig
cd docker/dig
emacs Dockerfile
The Dockerfile contents should look like this:
FROM tensorflow/tensorflow:latest-gpu-jupyter
RUN python3 -m pip install pandas
CMD ["bash", "-c", "source /etc/bash.bashrc && jupyter notebook --notebook-dir=/tf --ip 0.0.0.0 --no-browser --allow-root"]
docker build -t dig:tf-1 .
mkdir -p ~/notebooks/dig
docker run --gpus=all -it --rm \
-u $(id -u):$(id -g) \
-p 8888:8888 \
-v $HOME/notebooks/dig:/tf/dig \
dig:tf-1
To run the image without GPU support, just remove the --gpus=all option.
There are a bunch of nvidia-driver-XXX packages. You can see them with apt search nvidia-driver. I don’t know why I chose nvidia-driver-455.
For detailed information about installing Docker, see Install Docker Engine on Ubuntu.
For detailed information about installing the NVIDIA Container Toolkit, see NVIDIA Docker Installation Guide.
There might be a better way to install Python packages and make other changes to the default TensorFlow image. I opted for using a Dockerfile to create a new image based on the TensorFlow image.
Creating and then mounting the notebooks/dig directory ensures that your notebooks persist.
In general, it’s better to work without GPU support, and then running the container with GPU support if you have something that’s ready to run.
I put the Docker run command in a script, so I usually start my notebook like this: start-notebook or start-notebook-gpu.