python-3.x - 如何将 Tensorflow 2.0 SavedModel 转换为 TensorRT？

Question

我已经在 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: 没有名为“tensorflow.python.compiler.tensorrt”的模块

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

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

Answer 1

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

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