Build-TensorFlow-with-TensorRT

This repository contains helper Dockerfiles and build scripts to produce TensorFlow Python wheels built with NVIDIA TensorRT support for specific Python, CUDA, and cuDNN combinations.

Project Status

This project was created to facilitate building TensorFlow with TensorRT support on various Linux distributions. Currently, wheels have been successfully built and tested for:

• TensorFlow 2.13.0 & 2.13.1 - Complete (Ubuntu 18.04, 20.04, 22.04)
• TensorFlow 2.14.0 & 2.14.1 - Complete (Rocky Linux 8, Ubuntu 18.04, 20.04, 22.04)
• TensorFlow 2.15.0 & 2.15.1 - Python 3.9-3.11 complete (TODO: Add Ubuntu 18.04, 20.04, 22.04, and 24.04 support)
• TensorFlow 2.16.1 & 2.16.2 - Python 3.9-3.12 complete (TODO: Add Ubuntu 18.04, 20.04, 22.04, and 24.04 support; update Rocky Linux Dockerfiles)
• TensorFlow 2.17.0 & 2.17.1 - Python 3.9-3.12 complete (TODO: Add Ubuntu 18.04, 20.04, 22.04, and 24.04 support; update Rocky Linux Dockerfiles)

Build Wheels Availability - Wheels will be hosted soon on majorcadencetech.org (TODO: Upload and host wheels)

Next Steps

The following improvements are planned:

• Add Ubuntu 18.04, 20.04, 22.04, and Ubuntu 24.04 support for TensorFlow 2.15.x, 2.16.x, and 2.17.x
• Update Rocky Linux Dockerfiles for TensorFlow 2.16.x and 2.17.x to match current build standards
• Host pre-built wheels on majorcadencetech.org for easy access
• Expand ARM/Jetson support (future consideration)

Background

This project was born out of the work of converting Keras models to TensorRT models on an HPC running Rocky Linux 9. Ubuntu wheels are also provided. There is no ARM support (e.g., Jetson) at this time, although work on that may be planned.

Status: This is not feature complete. While I can verify that the wheels build successfully, comprehensive testing across all combinations is impractical. Community feedback and contributions are welcome.

The TensorRT version used is 8.6.1, the latest offered by NVIDIA that is compatible with the supported TensorFlow versions. TensorRT versions > 9 would require manual patching.

Prerequisites

• Linux host with Docker installed (no GPU is needed to build wheels, but one is needed to use them)
• CUDA Toolkit - Download from NVIDIA CUDA Downloads matching the version in your build configuration
• cuDNN - Manually download from NVIDIA cuDNN Downloads - ensure the version matches your build (check the Dockerfile for the exact version)
• TensorRT - Manually download from NVIDIA TensorRT Downloads - currently using version 8.6.1
• Sufficient disk space and memory for building TensorFlow (tens of GBs disk space, 16+ GB of RAM)
• Significant time (4-8 hours per build)

Important Notes on NVIDIA Libraries

• cuDNN and TensorRT must be manually downloaded from NVIDIA's website (they require an account and accept license terms)
• Verify version compatibility: The versions of CUDA, cuDNN, and TensorRT used during the build must match the versions you install on your deployment system
• These wheels do not bundle CUDA, cuDNN, or TensorRT libraries - you must install them separately on any system using these wheels
• Ensure all NVIDIA libraries are in your $LD_LIBRARY_PATH for proper runtime functionality

Usage

Step-by-Step Build Instructions

1. Choose the configuration matching your desired TensorFlow/Python/CUDA/cuDNN combination:

cd 2.16.2/python3.10/tf-cuda12.1-cudnn8.9-TRT-8.6.1/

2. Download required libraries (if not already installed):

   • Download cuDNN from NVIDIA cuDNN Downloads
   • Download TensorRT from NVIDIA TensorRT Downloads
   • Note the installation paths for use in the Dockerfile if needed

3. Build the Docker image:

docker build -t build_tf -f Dockerfile .

4. Enter the Docker container:

docker run --name TFBUILD -it build_tf /bin/bash

5. Start the TensorFlow build (inside the container):

./build_tf.sh

The build_tf.sh script executes the TensorFlow build. Edit this script if you need to change build configuration, flags, or dependencies. Any changes require rebuilding the Docker image.

6. Export the built wheel (from your host):

docker cp TFBUILD:/output ./

7. Verify the installation:

pip install <your-built-wheel-file>
python -c "
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert
print('TensorFlow version:', tf.__version__)
print('GPU devices:', tf.config.list_physical_devices('GPU'))
print('Build info:', tf.sysconfig.get_build_info()['cuda_version'])
"

Alternatively, run the included test script:

python test-tf.py

Customization

You can customize the build process by modifying files in the chosen configuration directory:

• Dockerfile - Modify base image, dependencies, or library versions
• build_tf.sh - Change Bazel build flags, TensorFlow configuration, or Python version
• requirements_lock_*.txt - Add or update Python package dependencies

GPU Compute Capability

These builds are configured for compute capability 8.6 (RTX 30 series and A6000). To change this for your GPU:

1. Find your GPU's compute capability on NVIDIA's GPU Computing Capability Chart
2. Edit the build_tf.sh script and update:
   • export TF_CUDA_COMPUTE_CAPABILITIES=8.6 (change to your GPU's capability)
   • --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=8.6 in the Bazel build command

Troubleshooting

Build Issues

• Docker build fails with CUDA errors: Ensure your host NVIDIA driver is compatible with the CUDA runtime version used by the image
• Out-of-memory errors: Increase Docker build resources (memory and disk). Try adding --jobs=x to the Bazel build command (where x is the number of CPU cores), which dramatically reduces memory usage at the cost of build time
• Build times are extremely long: This is normal. Building TensorFlow with TensorRT support can take 4-8 hours. Using --jobs flag appropriately can help

Runtime Issues

• TensorRT warnings (TF-TRT Warning: Could not find TensorRT):

  • Run python test-tf.py to diagnose library loading issues
  • Verify CUDA and cuDNN are installed and their .so files are in $LD_LIBRARY_PATH
  • Check that TensorRT library files (libcublas, libnvinfer, etc.) are in $LD_LIBRARY_PATH

• Symbol not found errors for .so libraries:

  • Ensure all NVIDIA libraries (CUDA, cuDNN, TensorRT) are in your $LD_LIBRARY_PATH
  • Example: export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/lib:$LD_LIBRARY_PATH

• GPU not detected (No GPU devices found):

  • Verify GPU is present: nvidia-smi
  • Check CUDA installation: ls /usr/local/cuda/lib64/libcuda.so*
  • Ensure CUDA libraries are in $LD_LIBRARY_PATH

• Version mismatches: If you get errors related to TensorRT, CUDA, or cuDNN versions:

  • Check your installed library versions using python test-tf.py
  • Ensure they match the versions specified in the Dockerfile and build_tf.sh
  • Reinstall matching library versions if necessary

Important Reminders

• These wheels do not bundle CUDA or cuDNN - Unlike pip install tensorflow[and-cuda], which includes CUDA and cuDNN, these wheels require you to install NVIDIA libraries separately
• TensorRT must be installed - Download and install TensorRT 8.6.1 on any system using these wheels
• Version consistency is critical - The CUDA, cuDNN, and TensorRT versions used during build must match what you install on deployment systems

Contributing

To contribute new configurations or improvements:

1. Add a new folder following the existing naming pattern for additional Python/TensorFlow/CUDA/cuDNN/TensorRT combinations
2. Include a Dockerfile and build_tf.sh script configured for your combination
3. Test the build to ensure it completes successfully
4. Submit a pull request with documentation of your testing

Acknowledgments

This project leverages open-source tools and NVIDIA's TensorRT library. Special thanks to the TensorFlow and NVIDIA communities for their documentation and support.

License

Please refer to the LICENSE file if included. This project builds TensorFlow under its Apache 2.0 license.