TensorFlow集成TensorRT环境配置

转自：AI科技大本营

本文前提是cuda和cudnn以及TensorRT已经安装完毕，具体详情，可以参考上一篇文章：

https://developer.nvidia-china.com/forum.php?mod=viewthread&tid=8767&extra=page%3D1

TensorRT下载地址：

https://developer.nvidia.com/nvidia-tensorrt-download 

TesnsoRT的介绍文档：

https://devblogs.nvidia.com/tensorrt-3-faster-tensorflow-inference/ 

TensorRT的开发者指南：

http://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html 

TensorRT的样例代码：

http://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#samples

当前的系统环境是：

• Ubuntu 16.04

• CUDA 9.0

• CUDNN 7.1

• TensorRT 4.0

 

1. 首先从TensorFlow的官方github上下载TensorFlow：

sudo apt-get install git

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

2.安装TensorFlow的依赖库：

sudo apt-get install 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文件

sudo pip install /tmp/tensorflow_pkg/tensorflow-1.2.0rc2-cp27-cp27mu-linux_x86_64.whl

直接安装即可。

 

如果你想测试一下效果，可以在这里下载一段测试代码：

https://developer.download.nvidia.com/devblogs/tftrt_sample.tar.xz

测试结果是这样的：

中国人工智能学会将于5月19日-20日举办“2018全球人工智能技术大会”，期间开展Google TensorFlow培训课程，详情点击“阅读全文”