您的位置:

Ubuntu安装TensorRT

一、下载安装包

首先从NVIDIA官网下载所需的TensorRT版本,并将其安装包保存到Ubuntu系统中的特定文件夹中。我们以TensorRT 7.2.1为例,打开Terminal,输入以下命令下载压缩包:

wget https://developer.download.nvidia.cn/compute/redist/tensorrt/tensorrt-7.2.1.6.linuxo-x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz

下载完成后,解压缩文件:

tar -xzvf tensorrt-7.2.1.6.linuxo-x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz

将解压缩的文件夹复制到/usr/src目录下:

sudo mv tensorrt-7.2.1.6 /usr/src/

二、安装依赖库

在安装TensorRT之前,需要确保系统已安装必需的依赖库,包括CUDA和cudnn等。请使用以下命令来安装:

sudo apt-get update
sudo apt-get install build-essential cmake git python3-dev python3-numpy python3-py python3 python-pytest python-pytest-xdist python-setuptools python-wheel libcudnn8 libcudnn8-dev libcurl4-openssl-dev libssl-dev wget

三、为TensorFlow安装TensorRT插件

要在TensorFlow中使用TensorRT,需先为TensorFlow安装对应的TensorRT插件。以下是具体步骤:

1、使用Git获取TensorRT代码:

git clone https://github.com/tensorflow/tensorrt

2、将TensorRT插件的源码文件夹“tensorflow”复制到TensorFlow根目录中的“tensorflow/contrib”目录下:

cp -r tensorflow /path/to/tensorflow/contrib/

3、回到TensorFlow根目录,运行bazel build进行编译:

cd /path/to/tensorflow
bazel build --config=cuda //tensorflow/contrib/tensorrt:trt_engine_op

4、使用以下命令在用户主目录下创建TensorRT的符号链接:

sudo ldconfig /usr/src/tensorrt-7.2.1.6/lib/

四、设置环境变量

为避免每次使用TensorRT都需要手动设置环境变量,我们可以设置它们以自动加载。将以下命令添加到您的Bashrc文件中:

echo 'export LD_LIBRARY_PATH=/usr/src/tensorrt-7.2.1.6/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
echo 'export PATH=/usr/src/tensorrt-7.2.1.6/bin:$PATH' >> ~/.bashrc

完成后,运行以下命令以使其生效:

source ~/.bashrc

五、测试TensorRT

为确保TensorRT已正确安装并在Ubuntu上运行,请打开Terminal,输入以下命令:

python3
>>> import tensorrt as trt
>>> print(trt.__version__)

如果命令行显示TensorRT的版本号,则表明TensorRT已安装并可供使用。