This is a short tutorial for installing Tensorflow 2.2 with Cuda 10.2, CuDNN 8 and TensorRT 7.1 on Ubuntu 18.04 with Python 3.6 and GCC 7.5. It works on my machine as of 23. 07. 2020 :) However, if you're reading this, the guide is probably already outdated.
If you have the exact same specification as I do, you can try your luck with my prebuilt pip wheel: https://mega.nz/file/MAozBYhI#8Tp7W8t-X5sBuWCVEviTN5quQnNNwaVMv90OAnEvGD0
Recently I got hold of Nvidia Xavier NX, which has JetPack 4.4 with Tensorflow 2.2 and TensortRT 7.1. I wanted to have the same environment on my desktop/server, but sadly, Tensorflow does not support CUDA 10.2, TensorRT 7.1 and CuDNN 8 by default in 2.2 and probably will not support it even in 2.3 (tensorflow/tensorflow#38194).
There are a few hiccups related to compiling Tensorflow with these specific library versions, so I put a short description of the process here. My guide mostly follows this awesome tutorial: https://gist.github.com/kmhofmann/e368a2ebba05f807fa1a90b3bf9a1e03, however it adds support for TensorRT 7.1 and CuDNN 8. Check that tutorial out if you need more information.
The first step is to install necessary dependencies. I won't detail installing CUDA, as that is out of scope, I therefore assume that you already have NVIDIA and NVIDIA machine learning APT repositories set up (if not, follow this guide: https://www.tensorflow.org/install/gpu).
# install CUDA 10.2, CuDNN 8, TensorRT 7.1 and Bazel 2.0.0
$ sudo apt install cuda-10-2 libcudnn-8 libcudnn8-dev libnvinfer-plugin-dev=7.1.3-1+cuda10.2 \
libnvinfer-plugin7=7.1.3-1+cuda10.2 libnvinfer-dev=7.1.3-1+cuda10.2 libnvinfer7=7.1.3-1+cuda10.2 \
build-essential openjdk-8-jdk python zip unzip curl
$ curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
$ echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
$ sudo apt update && sudo apt install bazel-2.0.0
I recommend to create a new virtual environment for building Tensorflow.
$ python3 -m venv venv
$ source venv/bin/activate
$ pip install -U pip six "numpy<1.19.0" wheel setuptools mock 'future>=0.17.1'
$ pip install -U keras_applications --no-deps
$ pip install -U keras_preprocessing --no-deps
Make sure that you install numpy below version 1.19.0, otherwise you're going to run into trouble during TF build.
First, download TF sources and checkout the 2.2 commit.
$ git clone https://github.com/tensorflow/tensorflow
$ cd tensorflow
$ git checkout v2.2.0
Since this version does not support CuDNN 8 and CUDA 10.2 out of the box, we need to cherry pick a commit that adds support for it:
# Add CUDA 10.2 + CuDNN 8 support
# run this in the root of the TF directory
$ git cherry-pick 255f590ab64e637f49288883013d35efa0633b35
$ git checkout 255f590ab64e637f49288883013d35efa0633b35 -- third_party/gpus/find_cuda_config.py.gz.base64
$ git add third_party/gpus/find_cuda_config.py.gz.base64
$ git cherry-pick --continue
And then we need to solve some additional compilation issues (tensorflow/tensorflow#40688):
# run this in the root of the TF directory
$ perl -pi.bak -e 's%, CompareUFunc%, (PyUFuncGenericFunction) CompareUFunc%g' tensorflow/python/lib/core/bfloat16.cc
This step might not be needed if you have numpy <1.19.0, but I for sure will not try to build TF again to check it.
Now all that's left to do is configure Tensorflow and build it. Make sure to enable CUDA during configuration and select the compute capability that you need (Xavier NX needs 7.2).
$ ./configure
Now you can finally run the build:
TMP=/tmp bazel build --config=opt -c opt //tensorflow/tools/pip_package:build_pip_package
Prepare for the build to take several hours. It is also possible that it will crash and/or segfault several times, in that case just run the above command again.
Once the build finishes, you can use the created binary to build a Python package with Tensorflow:
$ ./bazel-bin/tensorflow/tools/pip_package/build_pip_package package
You can find the resulting package in <tensorflow-root-dir>/package/tensorflow-2.2.0-cp36-cp36m-linux_x86_64.whl and install it with
$ pip install package/tensorflow-2.2.0-cp36-cp36m-linux_x86_64.whl