使用NVIDIA TensorRT推理引擎运行Tensorflow

时间:2016-12-14 12:09:12

标签: tensorflow nvidia tensorrt

我想使用NVIDIA TensorRT来运行我的Tensorflow模型。 Currenly,TensorRT支持Caffe prototxt网络描述符文件。

我无法找到将Tensorflow模型转换为Caffe模型的源代码。有没有解决方法?

2 个答案:

答案 0 :(得分:13)

TensorRT 3.0支持通过它的UFF(通用框架格式)导入/转换TensorFlow图。缺少某些层实现,需要通过IPlugin接口进行自定义实现。

以前的版本不支持TensorFlow模型/检查点的原生导入。

您还可以将层/网络描述导出为您自己的中间格式(例如文本文件),然后使用TensorRT C ++ API构建推理图。您必须单独导出卷积权重/偏差。确保注意重量格式 - TensorFlow使用NHWC,而TensorRT使用NCHW。对于权重,TF使用RSCK([filter_height,filter_width,input_depth,output_depth]),TensorRT使用KCRS。

有关张量格式的扩展讨论,请参阅本文: https://arxiv.org/abs/1410.0759

此链接也提供了有用的相关信息: https://www.tensorflow.org/versions/master/extend/tool_developers/

答案 1 :(得分:2)

由于新的TensorRT 3增加了对TensorFlow的支持,目前无需解决方法。