install_tensorrt.sh
Shebang and set options
The script starts with the shebang
#!/bin/bash
, indicating that it should be executed using the Bash shell.The
set -ex
command enables two options:-e
to exit the script immediately if any command fails, and-x
to print each command before executing it for debugging purposes.
Variable declarations
The script declares several variables with default values for TensorRT, CUDA, cuDNN, NCCL, and cuBLAS versions.
These variables are used throughout the script to specify the versions of the libraries to be installed.
Command-line argument parsing
The script uses a
for
loop to iterate over the command-line arguments passed to the script.It uses a
case
statement to check for specific argument patterns and updates the corresponding variables accordingly.This allows the user to override the default versions of the libraries by passing command-line arguments.
NVCC version check
The script uses the
nvcc --version
command to retrieve the version of the pre-installed CUDA compiler (NVCC).It compares the NVCC version with the
CUDA_VER
variable usinggrep
andhead
commands.If the versions don't match, the script prints an error message and exits with a non-zero status.
install_ubuntu_requirements
function
install_ubuntu_requirements
functionThis function is responsible for installing the required dependencies on Ubuntu.
It updates the package list, installs the NVIDIA CUDA keyring package, and removes any existing installations of libcudnn8, libnccl, libcublas, and cuda-nvrtc-dev.
It then installs the specified versions of libcudnn8, libnccl2, libcublas, and cuda-nvrtc-dev using
apt-get
.Finally, it cleans up the package cache and removes the package lists.
install_centos_requirements
function
install_centos_requirements
functionThis function is responsible for installing the required dependencies on CentOS.
It updates the package repositories, installs the EPEL repository, and removes any existing installations of libcudnn, libnccl, and libcublas.
It then installs the specified versions of libcudnn8, libnccl, and libcublas using
yum
.Finally, it cleans up the package cache.
install_tensorrt
function
install_tensorrt
functionThis function is responsible for installing TensorRT.
It determines the Python version using
python3 -c
and parses it to remove the dots.It sets the TensorRT CUDA version to "12.2".
If the
RELEASE_URL_TRT
variable is not set, it determines the appropriate download URL based on the target architecture and operating system.It downloads the TensorRT package using
wget
, extracts it to/usr/local/
, and renames the directory totensorrt
.It installs the TensorRT Python wheel package using
pip3
.Finally, it adds the TensorRT library path to the
LD_LIBRARY_PATH
environment variable by appending it to theENV
file.
OS detection and installation
The script uses
grep
andtr
commands to extract the OS ID from the/etc/os-release
file.It uses a
case
statement to determine the appropriate installation function based on the OS ID.If the OS is Ubuntu, it calls the
install_ubuntu_requirements
andinstall_tensorrt
functions.If the OS is CentOS, it calls the
install_centos_requirements
andinstall_tensorrt
functions.If the OS is not recognized, it prints an error message and exits with a non-zero status.
Overall, the install_tensorrt.sh
script is a Bash script that automates the installation of TensorRT and its dependencies on Ubuntu and CentOS systems.
It provides flexibility by allowing the user to specify the versions of the libraries through command-line arguments.
The script follows a modular approach by separating the installation steps into functions based on the OS. It also includes error handling and version checks to ensure compatibility and smooth execution.
Last updated