TensorRT是Nvdia推出的专门针对部署阶段进行加速的驱动,可用于加速模型推断,提供C++和Python接口,支持tensorflow和caffe框架,目前最新版本是4.0。
- 安装
我的环境是ubuntu Server 16 , anaconda3.6(官方推荐的python版本是2.7或3.5),CUDA9.0+cuDNN7.1(必要)。
首先下载TensorRT4.0的安装包,我选择的是基于cuda9.0的tar格式安装包。
解压文件:
tar -xzvf TensorRT-4.0.1.6.Ubuntu-16.04.4.x86_64-gnu.cuda-9.0.cudnn7.1.tar.gz
修改本地环境变量:
vim ~/.bashrc
添加 :
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:~/TensorRT-4.0.1.6/lib:/usr/local/cuda-9.0/lib64"
刷新生效:
source ~/.bashrc
然后安装python库:
cd ~/TensorRT-4.0.0.16/python
pip install tensorrt-4.0.0.16-cp35-cp35mu-linux_x86_64.whl
如果无法安装,请把文件名修改为:tensorrt-4.0.0.16-cp36-cp36mu-linux_x86_64.whl
用同样的方法安装UFF,graphsurgeon的python安装包
cd ~/TensorRT-4.0.1.6/uff
cd ~/TensorRT-4.0.1.6/graphsurgeon
import tensorrt as trt 测试python版本安装是否正确。由于是python3.6版本的,所以会无法导入某些包。我的处理办法是修改库目录的.so文件名:
cd ~/anaconda3/lib/python3.6/site-packages/tensorrt
搜索.so文件,将例如_nv_utils_bindings.cpython-35m-x86_64-linux-gnu.so修改为_nv_utils_bindings.cpython-36m-x86_64-linux-gnu.so即可解决问题。
最后测试C++接口:
cd ~/TensorRT-4.0.1.6/targets/x86_64-linux-gnu/samples/sampleMNIST
make CUDA_INSTALL_DIR=/usr/local/cuda
cd ~/TensorRT-4.0.0.16/bin
./sample_mnist
输出:
到此TensorRT安装成功。
官方文档:https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html