TensorFlow Installation

I have to admit that it is painful to install TensorFlow, especially when you need a specific version. That's why I took some notes.

My general experience is that the pip approach has many potential issues on its path, mostly due to the dependency version issues from NVIDIA's libraries. The conda approach is much easier as it bundles everything.

This is written on 2020-11-27. The applicability might change in the future.

GPU Driver

Follow the TensorFlow official guide to install GPU driver (ONLY the driver). This code snippet was from 2020-11-27 (versions might change later):

# Add NVIDIA package repositories
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.1.243-1_amd64.deb
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
sudo dpkg -i cuda-repo-ubuntu1804_10.1.243-1_amd64.deb
sudo apt-get update
wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
sudo apt install ./nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb
sudo apt-get update

# Install NVIDIA driver
sudo apt-get install --no-install-recommends nvidia-driver-450
# Reboot. Check that GPUs are visible using the command: nvidia-smi

Anaconda

Install Anaconda. Download it from the official website, then

chmod +x Anaconda3-2020.11-Linux-x86_64.sh
./Anaconda3-2020.11-Linux-x86_64.sh

TensorFlow

Install TensorFlow from Anaconda (bundled with proper CUDA, etc.). Follow the official guide. For example, to install a specific version (1.14.0) that works on GPU:

conda create -n tf-gpu tensorflow-gpu==1.14.0
conda activate tf-gpu

The installed conda environment will be under ~/anaconda3/envs/tf-gpu/.

Others

Install a specific version of numpy to avoid the error message when importing TensorFlow 1.14.0:

conda install "numpy<1.17"

Install a specific version of TensorFlow Probability (per greta's requirement):

conda install tensorflow-probability==0.7.0

Test

Test if TensorFlow runs on GPU:

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

with tf.device('/gpu:0'):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)

with tf.Session() as sess:
    print (sess.run(c))