0%

TensorRT 教程(一):安装TensorRT

TensorRT 是一个用于在 NVIDIA GPU 上进行高性能计算的 C++ 库。TensorRT 在运行时,加载一个已经训练好的神经网络,创建一个经过内部高度优化的引擎(engine),来进行快速计算。TensorRT 同时提供 C++ 和 Python API 接口。TensorRT 同时支持 Windows/Ubuntu/iOS 系统,本教程将基于 Ubuntu 进行讲解。

1. 安装依赖

TensorRT 具有以下依赖,请先安装:

  • CUDA 9.0, 10.0 ,或者 10.2
  • cuDNN 7.6.5
  • Python 2 / Python 3 (可选)

2. 下载TensorRT

根据 TensorRT 官网介绍,使用压缩包安装可方便后续开发,当然也可选择 Debian/NPM 安装方式。本文将按照 tar 压缩包方式进行安装。按照下面流程即可快速安装,记得先注册一个 NVIDIA 账号:

  1. 进入官网:https://developer.nvidia.com/tensorrt
  2. 点击 Download Now,然后登录NVIDIA账号
  3. 选择需要安装的TensorRT版本(此处选择 7.0)
  4. 同意软件安装许可
  5. 点击软件包下载链接,随后开始下载

3. 安装 TensorRT

3.1 安装 TensorRT 源码

将下载好的压缩包,放在一个路径下,比如在我的电脑:/home/TensorRT-7。然后将 TensorRT 的 lib 文件路径加入系统环境变量:

1
2
3
sudo gedit ~/.bashrc`  # sudo vi ~/.bashrc
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<TensorRT-${version}/lib>
source ~/.bashrc

3.2 安装 Python 接口

如果需要用 Python API 进行编程,按照如下进行:

1
2
3
cd TensorRT-${version}/python
sudo pip2 install tensorrt-*-cp27-none-linux_x86_64.whl # python2
# sudo pip3 install tensorrt-*-cp3x-none-linux_x86_64.whl # python3

3.3 安装UFF转换库

如果你想将 TensorFlow 训练后的网络通过 UFF 编码方式转换给 TensorRT 进行运行:

1
2
3
cd TensorRT-${version}/uff
sudo pip2 install uff-0.6.5-py2.py3-none-any.whl # python2
# sudo pip3 install uff-0.6.5-py2.py3-none-any.whl # python3

3.4 安装 graphsurgeon

graphsurgeon 是对UFF编码网络进行定制化操作的库,比如插入或删除神经网络某一层layer:

1
2
3
cd TensorRT-${version}/graphsurgeon
sudo pip2 install graphsurgeon-0.4.1-py2.py3-none-any.whl # python2
# sudo pip3 install graphsurgeon-0.4.1-py2.py3-none-any.whl # python3

4. 验证安装

  1. 查看TensorRT的安装目录下文件是否齐全,可使用命令tree -d,会看到包含以下文件夹:libincludedata
  2. 运行例子sampleMNIST:
    1
    2
    3
    cd <TensorRT root directory>/samples/sampleMNIST
    make
    ./sample_mnist
    如果上述命令无ERROR输出,则证明安装成功。

至此,TensorRT安装成功!