tensorflow2.0 h5 模型转化 pb

原创

已于 2022-03-31 15:38:24 修改 · 2.6k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#python

于 2022-03-29 12:10:11 首次发布

本文介绍了如何将 TensorFlow 2.0 的 H5 模型转换为 PB 格式，并提供了转换命令。此外，还提及了模型进一步转换为 ONNX 格式的过程，利用 tf2onnx 工具完成这一操作。


import tensorflow as tf

import os
from tensorflow.python.tools import freeze_graph
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

import sys

def h5_to_pb(h5_save_path):
    model = tf.keras.models.load_model(h5_save_path, compile=False)
    model.summary()
    full_model = tf.function(lambda Input: model(Input))
    full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))

    # Get frozen ConcreteFunction
    frozen_func = convert_variables_to_constants_v2(full_model)
    frozen_func.graph.as_graph_def()

    layers = [op.name for op in frozen_func.graph.get_operations()]
    print("-" * 50)
    print("Frozen model layers: ")
    for layer in layers:
        print(layer)

    print("-" * 50)
    print("Frozen model inputs: ")
    print(frozen_func.inputs)
    print("Frozen model outputs: ")
    print(frozen_func.outputs)

    # Save frozen graph from frozen ConcreteFunction