使用 NVIDIA TensorRT 推理引擎运行 TensorFlow

Posted 2023-04-15

【中文标题】使用 NVIDIA TensorRT 推理引擎运行 TensorFlow

【英文标题】：Run Tensorflow with NVIDIA TensorRT Inference Engine

【发布时间】：2017-04-29 18:20:56

【问题描述】：

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

我找不到将 TensorFlow 模型转换为 Caffe 模型的源代码。有什么解决方法吗？

【参考方案1】：

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/

【讨论】：

我读到 tensorflow 1.7+ 支持 tensorRT。这意味着您需要合适的硬件

【参考方案2】：

由于新的 TensorRT 3 添加了对 TensorFlow 的支持，目前不需要任何变通方法。