Установка TensorFlow на Ubuntu 18.04/CUDA/TensorRT

30.06.202005.07.2020 Александр Кручинин Tensorflow, tensorrt

Эта статья была написана по причине того, что текущая официальная страница установки для голой Ubuntu 18.04 GPU не делает работоспособной конфигурацию TensorFlow с CUDA и TensorRT. Поскольку мне надо было работать с frozen graph, который не поддерживается интерфейсом TensorFlow 2.x для оптимизации TensorRT, то ставилась версия 1.15 Tensorflow. Итак, приступим.

1. Начальные обновления плюс виртуальная среда.

Смотрим версию pip3. Она очень старая и не позволяет установить нужную версию TensorFlow:

ubuntu@ubuntu-System-Product-Name:~$ pip3 --version
pip 9.0.1 from /usr/lib/python3/dist-packages (python 3.6)

Не рекомендуется обновлять pip, поскольку на него завязана операционная система, поэтому лучше создать виртуальную среду и установить всё, что нужно там. Устанавливаем нужные компоненты:

sudo apt update

sudo apt install python3-dev python3-pip python3-venv

sudo apt-get update

sudo apt install python3-pip

И создаём виртуальную среду:

python3 -m venv --system-site-packages ./venv

Переходим в среду:

source ./venv/bin/activate

Должен появиться слева текст venv:

(venv) ubuntu@ubuntu-System-Product-Name:~$

Выйти из среды всегда можно командой deactivate, но пока нам этого не нужно, поскольку обновляем pip:

pip install --upgrade pip

Смотрим версию:

(venv) ubuntu@ubuntu-System-Product-Name:~$ pip3 --version
pip 20.1.1 from /var/www/ubuntu/t/venv2/lib/python3.6/site-packages/pip (python 3.6)

То, что нужно.

2. Устанавливаем TensorFlow

pip3 install tensorflow-gpu==1.15

Идёт скачивание и установка TensorFlow. Проблем здесь обычно не возникает.

3. NVIDIA драйвер

Необходимо скачать и установить драйвер NVIDIA. Можно сделать так, как указано на сайте TensorFlow, хотя всё равно установится более новая версия:

sudo apt-get install --no-install-recommends nvidia-driver-430

Перезагружаем и проверяем командой nvidia-smi:

Драйвер установлен

4. Устанавливаем CUDA 10 и CUDNN

Для этого скачиваем его отсюда:

https://developer.nvidia.com/cuda-10.0-download-archive?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=1804&target_type=deblocal

А затем выполняем команды:

sudo dpkg -i cuda-repo-ubuntu1804-10-0-local-10.0.130-410.48_1.0-1_amd64.deb

sudo apt-key add /var/cuda-repo-10-0-local-10.0.130-410.48/7fa2af80.pub

sudo apt-get update

sudo apt-get install cuda

Если у вас стоит более новая версия CUDA, и поэтому старая не ставится, то просто выполните эту команду

sudo apt-get install cuda-10-0

Далее скачиваем его с сайта NVIDIA, тут вам придётся зарегистрироваться. Скачиваем тут (https://developer.nvidia.com/cudnn)и распаковываем затем архивы так:

sudo dpkg -i libcudnn7_7.6.5.32-1+cuda10.0_amd64.deb

sudo dpkg -i libcudnn7-dev_7.6.5.32-1+cuda10.0_amd64.deb

Проверяем работоспособность TensorFlow GPU

python3 test.py

Где тестовый файл такой:

import tensorflow as tf

tf.config.list_physical_devices('GPU')

Результат у меня такой:

2020-07-05 15:56:13.343641: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1304] Created TensorFlow device (/device:GPU:0 with 7184 MB memory) -> physical GPU (device: 0, name: GeForce RTX 2080 SUPER, pci bus id: 0000:01:00.0, compute capability: 7.5)
Found GPU at: /device:GPU:0

5. Устанавливаем TensorRT

Качаем TensorRT версии 5.1.5 deb файл отсюда (https://developer.nvidia.com/nvidia-tensorrt-download) и затем распаковываем его:

sudo dpkg -i nv-tensorrt-repo-ubuntu1804-cuda10.0-trt5.1.5.0-ga-20190427_1-1_amd64.deb
sudo apt-key add /var/nv-tensorrt-repo-cuda10.0-trt5.1.5.0-ga-20190427/7fa2af80.pub
sudo apt-get update
sudo apt-get install libnvinfer5=5.1.5-1+cuda10.0
sudo apt-get install libnvinfer-dev=5.1.5-1+cuda10.0
sudo apt-mark hold libnvinfer5 libnvinfer-dev

TensorRT установлен. Проверяем, конвертируется ли модель (я конвертировал модель детектирования объектов).

from tensorflow.python.compiler.tensorrt import trt
import tensorflow as tf
with tf.Session() as sess:
    # First deserialize your frozen graph:
    with tf.gfile.GFile("/var/www/ubuntu/t/in/saved_model.pb", 'rb') as f:
        frozen_graph = tf.GraphDef()
        frozen_graph.ParseFromString(f.read())
    # Now you can create a TensorRT inference graph from your
    # frozen graph:
    trt_graph = trt.create_inference_graph(
        input_graph_def=frozen_graph,
    outputs=["detection_boxes","detection_scores","detection_classes","num_detections"],
    max_batch_size=1,
        max_workspace_size_bytes=500000000,
    precision_mode='FP16',
    minimum_segment_size=50)
    #is_dynamic_op=True)
    trt_engine_opts = len([1 for n in trt_graph.node if str(n.op) == 'TRTEngineOp'])
    print("trt_engine_opts = {}".format(trt_engine_opts))
    with open('trt_graph2.pb', 'wb') as f:
        f.write(trt_graph.SerializeToString())