TensorFlow集成TensorRT环境配置
本文前提是cuda和cudnn以及TensorRT已经安装完毕,具体详情,可以参考上一篇文章:
https://developer.nvidia-china.com/forum.php?mod=viewthread&tid=8767&extra=page%3D1
当前的系统环境是:
Ubuntu 16.04
CUDA 9.0
CUDNN 7.1
TensorRT 4.0
1. 首先从TensorFlow的官方github上下载TensorFlow:
sudo apt-getinstall git
git clone https://github.com/tensorflow/tensorflow
2.安装TensorFlow的依赖库:
sudo apt-getinstall python-numpyswigpython-dev python-wheel
3.运行下载下来的TensorFlow中的./configure文件来配置环境:
这里有几个地方需要注意,在选择是否支持cuda,cudnn和TensorRT的时候要记得选择yes,并根据您的安装环境,配置好路径,否则后面会报错。
特别是cuda,cudnn和TensorRT的版本的时候,一定要仔细,否则后面会浪费很多时间。
当前我的系统配置的是CUDA 9.0 + CuDNN7.1 + TensorRT4.0
其他的配置看自己的需求。
4.通过bazel来编译pip的安装包,然后通过pip安装
这里的—config=opt 后面有一个空格
bazel build --config=cuda --config=opt//tensorflow/tools/pip_package:build_pip_package
bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/
然后在tmp下面可以看到pip安装所需要.whl文件
sudopip install /tmp/tensorflow_pkg/tensorflow-1.2.0rc2-cp27-cp27mu-linux_x86_64.whl
直接安装即可。
如果你想测试一下效果,可以在这里下载一段测试代码:
https://developer.download.nvidia.com/devblogs/tftrt_sample.tar.xz
测试结果是这样的: