搞TensorRT的安装快一个星期了,在TensorFlow、NVIDIA开发者的帮助下最终确定如下可行环境:
System: Ubuntu 18.04
TensorRT version: 3.0.4 (for Ubuntu 14.04)
CUDA version: 9.0(run或者deb)
cuDNN version: 7.1.3 (run或者deb)
TensorFlow version: 1.8.0
GPU: GTX 1080
安装CUDA 9.0/CuDNN 7.1.3
这个比较随意,无论是deb还是run文件,通吃。具体安装过程就省略了。下载TensorRT
注意,官方给出的建议是:TensorRT 3.0.4比较稳定,和TensorFlow兼容较好,而且,Ubuntu 14.04版的TensorRT可以通吃14.04、16.04、18.04
所以就用Ubuntu 14.04版的TensorRT 3.0.4就包你不会出错
下载倒数第二个
安装TensorRT
解压后,将你的TensorRT添加到bashrc中:$ vim ~/.bashrc export LD_LIBRARY_PATH=/home/micro/TensorRT 3.0.4/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
安装TensorFlow
TensorFlow的安装可以通过pip或者源码编译安装,注意源码安装需要指定TensorRT路径:... Do you wish to build TensorFlow with CUDA support? [y/N]: y CUDA support will be enabled for TensorFlow. Please specify the CUDA SDK version you want to use, e.g. 7.0. [Leave empty to default to CUDA 9.0]: 9.0 Please specify the location where CUDA 9.0 toolkit is installed. Refer to README.md for more details. [Default is /usr/local/cuda]: Please specify the cuDNN version you want to use. [Leave empty to default to cuDNN 7.0]: 7.1 Please specify the location where cuDNN 7 library is installed. Refer to README.md for more details. [Default is /usr/local/cuda]: Do you wish to build TensorFlow with TensorRT support? [y/N]: y TensorRT support will be enabled for TensorFlow. Please specify the location where TensorRT is installed. [Default is /usr/lib/x86_64-linux-gnu]:/home/micro/TensorRT-4.0.0.3 ...
- 测试
依旧通过tftrt文件中的run_all.sh进行测试,但是,如果你的GPU内存较少,可能会报“OOM”,可以通过一次运行一种精度模式以及较少batch size的方法较少内存占用:
$ python tftrt_sample.py --INT8 \
--num_loops 10 \
--topN 5 \
--batch_size 2 \
--workspace_size 2048 \
--log_file log.txt \
--network resnet_v1_50_frozen.pb \
--input_node input \
--output_nodes resnet_v1_50/predictions/Reshape_1 \
--img_size 224 \
--img_file grace_hopper.jpg