Install TensorFlow GPU Docker images

Nvidia provides "Nvidia Container Toolkit" which enable docker images to access GPU. The installation of GPU tensorflow is complicated. The official site recommends using Docker.

requirements

Docker: including and after 19.03

NVIDIA driver

internet

Install Nvidia Container Toolkit

According to:

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker

Take CentOS 7 as examples.

1. Setup stable repo:

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
   && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo

2. Install nvidia-docker2:

sudo yum clean expire-cache
sudo yum install -y nvidia-docker2

# Restart docker
sudo systemctl restart docker 

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

If return information of you GPU-drivers. The install is finished.

For other Linux distribution, please reference to the Nvidia guide.

Pull TensorFlow images

If you have the latest Nidia driver. You can pull the latest GPU version of tensorflow

docker run --gpus all -it --rm tensorflow/tensorflow:latest-gpu-jupyter \
   python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

Otherwise, if you are using old version Nvidia drive. Go to DockerHub search older tags.

If the command above return a tensorsum normally, your install is finished.

Update image

Enter into the container

docker run --gpus all -it  tensorflow/tensorflow:<Your version> bash

Install chemistry packages

pip install pyscf
pip install rdkit-pypi
pip install pandas
pip install seaborn
# ASE

Leave the container by typing exit . Check the container id by typing docker ps -a. Find the ID of your latest container.

And commit the change of the container.

docker commit -c 'CMD ["jupyter", "notebook", "--notebook-dir=/tf", "--ip", "0.0.0.0", "--no-browser", "--allow-root"]' <CONTAINER ID>  updatetfchem/tfchem:v1

updatetfchem/tfchem:v1 is the updated image name and version.

Add CMD to modify the default command of container. Type

docker run --gpus all -u $(id -u)  -it -p 1234:8888 updatetfchem/tfchem:v1

The inner port 8888 is redirected to port 1234. You can used the port 1234 to access your Jupyter notebook.