如何将 Tensorflow 2.0 SavedModel 转换为 TensorRT？答案

【问题标题】：How to convert Tensorflow 2.0 SavedModel to TensorRT?如何将 Tensorflow 2.0 SavedModel 转换为 TensorRT？
【发布时间】：2020-03-09 19:47:11
【问题描述】：

我已经在 Tensorflow 2.0 中训练了一个模型，并且正在尝试改进迁移到生产时的预测时间（在支持 GPU 的服务器上）。在 Tensorflow 1.x 中，我能够通过使用冻结图来获得预测加速，但从 Tensorflow 2 开始，这已被弃用。通过阅读 Nvidia 对 TensorRT 的描述，他们建议使用 TensorRT 与单独使用 Tensorflow 相比，可以将推理速度提高 7 倍。来源：

TensorFlow 2.0 with Tighter TensorRT Integration Now Available

我已经训练了我的模型并使用 Tensorflow 的 SavedModel 格式将其保存到 .h5 文件中。现在我按照 nvidia 的文档来优化模型以使用 tensorrt 进行推理：TF-TRT 2.0 Workflow With A SavedModel。

当我跑步时：

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

我收到错误：ModuleNotFoundError: No module named 'tensorflow.python.compiler.tensorrt'

他们在这里举了另一个 Tensorflow 2.0 的例子：Examples。但是，他们尝试导入与上面相同的模块，我得到了同样的错误。

谁能建议如何使用 TensorRT 优化我的模型？

【问题讨论】：

您使用的是什么版本的 TF，您是如何构建/安装它的？
@PooyaDavoodi - 我已经解决了我的问题（请参阅下面的答案）。您的问题是正确的，因为这是 TF 版本的问题 - 谢谢

标签： python-3.x tensorflow nvidia tensorflow2.0 tensorrt

【解决方案1】：

我已经解决了这个问题。问题是我在本地 Windows 机器上测试代码，而不是在支持 gpu 的 AWS EC2 实例上。

似乎 tensorflow.python.compiler.tensorrt 包含在 tensorflow-gpu 中，但不包含在标准 tensorflow 中。为了使用 TensorRT 转换 SavedModel 实例，您需要使用带有 tensorflow-gpu 的机器。（我知道运行模型需要这样做，但没有意识到转换模型需要这样做。）

【讨论】：