TensorFlow集成TensorRT环境配置

2018-04-26  本文已影响1296人  AI科技大本营

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

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

TensorRT下载地址

TesnsoRT的介绍文档

TensorRT的开发者指南

TensorRT的样例代码

当前的系统环境是:

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

测试结果是这样的:

上一篇下一篇

猜你喜欢

热点阅读