TensorFlow集成TensorRT环境配置

本文前提是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-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

测试结果是这样的:

AI科技大本营

公众号ID:rgznai100

  • 发表于:
  • 原文链接:http://kuaibao.qq.com/s/20180424A175AG00?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。

扫码关注云+社区

领取腾讯云代金券