將 Tensorflow Frozen Graph 轉換為 UFF 以進行 TensorRT 推理

Question

我想在 TensorRT (C++) 中使用 tensorflow 模型，首先需要將 .pb 格式轉換為 .uff 格式。 執行此代碼時：

uff_model = uff.from_tensorflow_frozen_model(
  frozen_file="my_tf_model.pb",
  output_nodes=["output_layer"],
  output_filename="tmp.uff",
  debug_mode=True,
)

我收到此錯誤消息：

Traceback (most recent call last):
  File "/home/jodo/ILONA/object-detection/ssd/src/uff_converter/uff_converter.py", line 29, in <module>
    text=False,
  File "/home/jodo/miniconda3/envs/uff_converter/lib/python3.7/site-packages/uff/converters/tensorflow/conversion_helpers.py", line 228, in from_tensorflow_frozen_model
    graphdef.ParseFromString(frozen_pb.read())
google.protobuf.message.DecodeError: Error parsing message

導出的圖 (my_tf_model.pb) 使用 tensorflow 2.0.0 進行訓練和保存。 但是對於 uff 轉換器，我必須使用 tensorflow 1.15.0。 這可能是一個問題還是 .pb 文件應該向下兼容？

更新：使用與使用 uff 轉換器 (1.15.0) 相同版本訓練的模型進行測試，但仍然存在相同的錯誤。

Answer 1

回答我自己的問題：我的 .pb 文件不是凍結圖，而是 SavedModel 格式的一部分

要解決此問題 1) 將其轉換為凍結圖並使用凍結圖：

python -m tensorflow.python.tools.freeze_graph --input_saved_model_dir SAVED_MODEL_DIR

然后uff.from_tensorflow_frozen_model()應該可以工作。