Skip to content

TensorFlow Installation

In this article

TensorFlow is an open-source platform for machine learning and deep neural network training, developed by Google. With TensorFlow, you can create, train, and utilize complex neural networks using a user-friendly API in Python.

This platform supports both traditional machine learning algorithms, such as linear and logistic regression, as well as more advanced architectures, including convolutional neural networks (Convolutional Neural Networks, CNN) and recurrent neural networks (Recurrent Neural Networks, RNN).

TensorFlow Installation on Linux

This instruction is suitable for the following operating systems: Ubuntu 22.04, and verified for Python versions: Python 3.10.

Note

If you plan to use GPU acceleration, please install NVIDIA drivers and CUDA according to this instruction.

  1. Install Python:

    sudo apt install python3.10

    In Ubuntu 22.04, this version is installed by default, so we do not recommend installing a newer version.

  2. Installing system libraries for 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. Create a virtual environment for Python:

    python3 -m venv venv

  4. Activate the virtual environment:

    source venv/bin/activate

    After successful activation, the prompt will include the name of the virtual environment in parentheses:

    (venv) user@49069:~$

    Note

    You can create as many virtual environments as you like and install different libraries (including simultaneously, but sometimes this may cause conflicts).

  5. Installing bindings for TensorRT:

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

    To verify the installation, we can run Python in the console and execute the following script:

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

    Upon successful installation, the output will be:

    (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. Installing TensorRT version 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. Installing the TensorFlow library with GPU support:

    pip install tensorflow[and-cuda]

  8. Exit the virtual environment:

    deactivate

  9. Create a script to run 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

Running TensorFlow

To run TensorFlow in a virtual environment with specified library variables from the root directory of your home directory, use the command:

. tensorflow.sh

Verifying TensorFlow Installation

To verify the functionality and GPU support of the library, launch the following Python program in the console:

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

Upon successful installation, you will receive an output similar to this, displaying your GPU usage.

(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')]
>>>

Additionally, you can download and run the test Python script for training a neural network 2c_nn_mnist_customtrain.py from this GitHub repository.

Upon completion of its execution, you should receive an output similar to this:

(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%)