29th December, 2021

Easy dev setup of TensorFlow on Ubuntu

I got a new laptop with a GPU and decided to re-start my deep learning development learning path with tensorflow (tf) from scratch. To install it and get it running well was a more involved process than I thought.

To begin with, confirm that your system has a GPU which is supported. There is a list at https://developer.nvidia.com/cuda-gpus (most consumer GPUs will be under CUDA-Enabled GeForce and TITAN Products) but it's not comprehensive (I have a 1660 GTX which is not in the list). As long as you have something in the same range as the GPUs mentioned in the link, you should be good.

Also, install the latest LTS version of ubuntu. I initially installed Ubuntu 20.10 (which is the latest version as of this note) and had to wipe it and install 20.04 LTS to get all the components working.

The official page suggests two options docker and pip. Additionally, there are installation options using conda:

Using docker

For me, this did not work out of the box and it was too much to figure out what was going wrong on two fronts - the tensorflow and docker container.

This option might work for you if are already adept at using docker. I would avoid otherwise. Its definitely not 'the easiest way to set up GPU support' as the website says.

Using conda

Next, I tried installing using conda because thats what I use in most of my development work.

The official anaconda installation at https://anaconda.org/anaconda/tensorflow-gpu is seamless in installing all the required components but its also very old (v2.4.2 for linux). I installed this initially and found a lot of tutorials at the tensorflow website were breaking. This might be a good option for windows users though.

The other installation from conda-forge https://anaconda.org/conda-forge/tensorflow/badges might be better (v2.6.2) but still not at the latest version (v2.7.0).

Using pip

This is the easiest and best option currently IMO. The directions at https://www.tensorflow.org/install/pip are pretty good.

  • CUDA Setup
  • Pip / virtualenv
    • Install pip3/virtualenv
      • > sudo apt update
      • > sudo apt install python3-dev python3-pip python3-venv
    • Create and activate your virtualenv
      • > python3 -m venv --system-site-packages ./venv
      • > source ./venv/bin/activate # sh, bash, or zsh
    • Upgrade pip
      • (venv) > pip install --upgrade pip
    • Insall Tensorflow
      • (venv) > pip install --upgrade tensorflow
    • Done!