Перейти к содержанию

Установка TensorFlow

В этой статье

TensorFlow — это открытая платформа для машинного обучения и обучения глубоких нейронных сетей, созданная компанией Google. С помощью TensorFlow можно создавать, обучать и использовать сложные нейронные сети, используя удобный API на языке Python.

Платформа поддерживает как классические алгоритмы машинного обучения, такие как линейная и логистическая регрессии, так и более сложные архитектуры, включая сверточные нейронные сети (Convolutional Neural Networks, CNN) и рекуррентные нейронные сети (Recurrent Neural Networks, RNN).

Установка Tensorflow под Linux

Инструкция подходит для следующих операционных систем: Ubuntu 22.04 и проверена следующих версиях Python: Python 3.10.

Внимание

Если вы планируете использовать GPU-ускорение вычислений, то установите драйвера на видеокарту NVIDIA и CUDA по данной инструкции.

  1. Устанавливаем Python: 

    sudo apt install python3.10

    В Ubuntu 22.04 данная версия установлена по умолчанию, поэтому мы не рекомендуем ставить версию новее.

  2. Устанавливаем системные библиотеки TensorRT:

    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get install tensorrt

  3. Создаем виртуальное окружение:

    python3 -m venv venv

  4. Активируем в данное виртуальное окружение:

    source venv/bin/activate

    При успешной активации в строке приглашения появится имя виртуального окружения в скобках:

    (venv) user@49069:~$

    Примечание

    Вы можете создавать сколько угодно таких окружений и ставить в них разные библиотеки (в том числе и совместно, но иногда это может привести к конфликту).

  5. Устанавливаем биндинги для TensorRT:

    python3 -m pip install wheel
python3 -m pip install --pre --upgrade tensorrt

    Проверить установку библиотек можно вызвав консоль Python (задав в командной строке python) и набрав там следующую программу:

    import tensorrt
print(tensorrt.__version__)
assert tensorrt.Builder(tensorrt.Logger())

    При успешной установке получаем следующий вывод:

    (tensorflow) user@49069:~/gpu$ python
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorrt
>>> print(tensorrt.__version__)
10.0.0b6
>>> assert tensorrt.Builder(tensorrt.Logger())
>>> import tensorrt_lean as trt
>>> print(trt.__version__)
10.0.0b6
>>> assert trt.Runtime(trt.Logger())
>>> import tensorrt_dispatch as trt
>>> print(trt.__version__)
10.0.0b6
>>> assert trt.Runtime(trt.Logger())
>>>

  6. Устанавливаем TensorRT версии 8.6.1:

    wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz

  7. Устанавливаем библиотеку TensorFlow с поддержкой GPU:

    pip install tensorflow[and-cuda]

  8. Выходим из виртуального окружения:

    deactivate

  9. Создаем скрипт для запуска TensorFlow:

    echo '#!/bin/bash' > tensorflow.sh
echo 'source venv/bin/activate' >> tensorflow.sh
echo 'export CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> tensorflow.sh
echo 'export LD_LIBRARY_PATH=$CUDNN_PATH/lib:/usr/local/cuda/lib64' >> tensorflow.sh
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/TensorRT-8.6.1.6/lib' >> tensorflow.sh
chmod +x tensorflow.sh

Запуск TensorFlow

Для запуска TensorFlow в виртуальном окружении python с прописанными переменными библиотек из корня домашней директории используем команду:

```bash
. tensorflow.sh
```

Проверка установки Tensorflow

Для проверки работы библиотеки и наличия поддержки GPU, необходимо запустить в консоли Python следующую программу:

import tensorflow as tf
print(tf.reduce_sum(tf.random.normal([1000, 1000])))
print(tf.config.list_physical_devices('GPU'))

При успешной установке вы получите вывод, похожий на этот, в котором будет отображено использование вашего GPU:

(venv) user1@49069:~$ python
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
nt(tf.config.list_physical_devices('GPU'))
2024-04-22 23:39:54.472502: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
>>> print(tf.reduce_sum(tf.random.normal([1000, 1000])))
2024-04-22 23:39:55.810888: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1928] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22282 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:07:00.0, compute capability: 8.9
tf.Tensor(332.5041, shape=(), dtype=float32)
>>> print(tf.config.list_physical_devices('GPU'))
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
>>>

Также вы можете скачать и запустить тестовый Python-скрипт тренировки нейросети 2c_nn_mnist_customtrain.py из этого репозитория на GitHub.

По итогу его работы вы должны получить вывод, похожий на этот:

(venv) user1@49069:~$ python ./2c_nn_mnist_customtrain.py
2024-04-23 10:35:25.832754: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-04-23 10:35:28.404381: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1928] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22282 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:07:00.0, compute capability: 8.9
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1713861338.172126   21279 service.cc:145] XLA service 0x7f36a3bd2ec0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1713861338.172187   21279 service.cc:153]   StreamExecutor device (0): NVIDIA GeForce RTX 4090, Compute Capability 8.9
2024-04-23 10:35:38.180648: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-04-23 10:35:38.207205: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:465] Loaded cuDNN version 8902
I0000 00:00:1713861338.355523   21279 device_compiler.h:188] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9022 time: 14.9 test-acc: 0.932 (error: 6.75%)
1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9474 time: 4.74 test-acc: 0.953 (error: 4.71%)
1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9628 time: 4.83 test-acc: 0.963 (error: 3.67%)
1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9734 time: 4.86 test-acc: 0.968 (error: 3.25%)
1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9798 time: 4.9 test-acc: 0.97 (error: 3.04%)