在Jetson Nano上为TensorRT转换TF 2.0保存的模型

时间:2019-11-14 19:26:03

标签: tensorflow tensorflow2.0 tensorrt nvidia-jetson nvidia-jetson-nano

我正在尝试在Jetson Nano上将TF 2.0 save_model转换为tensorRT。

模型已保存在TF 2.0.0中。 nano具有Jetpack 4.2.2 w / TensorRT __和Tensorflow 1.14(这是Jetson的最新Tensorflow版本)。

我一直遵循here中的说明,该说明描述了如何将TF 2.0.0 saved_model转换为TensorRT。

下面是我的代码:

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
tf.enable_eager_execution()

converter = trt.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir)
converter.convert()
converter.save(output_saved_model_dir)

saved_model_loaded = tf.saved_model.load(
    output_saved_model_dir, tags=[tag_constants.SERVING])
graph_func = saved_model_loaded.signatures[
    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
frozen_func = convert_to_constants.convert_variables_to_constants_v2(
    graph_func)
def wrap_func(*args, **kwargs):
    # Assumes frozen_func has one output tensor
    return frozen_func(*args, **kwargs)[0]
output = wrap_func(input_data).numpy()

似乎开始成功转换。但是,当到达KeyError: 'serving_default'行时出现convert_to_tensor错误。我的完整打印输出在下面找到here(对于SO来说太长了),但是下面显示了python traceback。我该如何解决?

谢谢!

打印输出摘要(完整的打印输出here):

Traceback (most recent call last):
  File "tst.py", line 38, in <module>
    convert_savedmodel()
  File "tst.py", line 24, in convert_savedmodel
    converter.convert()
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/compiler/tensorrt/trt_convert.py", line 956, in convert
    func = self._saved_model.signatures[self._input_saved_model_signature_key]
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/saved_model/signature_serialization.py", line 196, in __getitem__
    return self._signatures[key]
KeyError: 'serving_default'

1 个答案:

答案 0 :(得分:1)

我可以在您的实验中看到两个问题:

  • 您正在安装TF 1.14的同时使用TF-TRT 2.0 API。不支持。如果您的系统上安装了TF 1.14,则需要使用TF-TRT 1.x API。

  • 根据https://www.tensorflow.org/guide/versions

  • ,保存在TF2.0中的TF模型与TF1.14不兼容。

如果您只能访问TF1.14,建议在应用TF-TRT之前重新生成TF1.14中的图形并将模型保存在那里,然后使用TF-TRT 1.x API。