TensorFlow集成TensorRT环境配置

star2017 1年前 ⋅ 1421 阅读


本文前提是cuda和cudnn以及TensorRT已经安装完毕。


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


测试结果是这样的:


更多内容请访问:IT源点

相关文章推荐

全部评论: 0

    我有话说: