如何使用TF2将加载的h5模型正确保存到pb

Flo*_*Man 4 protocol-buffers keras tensorflow keras-2 tensorflow2.0

我加载了一个已保存的 h5 模型并希望将该模型保存为 pb。在训练期间使用tf.keras.callbacks.ModelCheckpoint回调函数保存模型。

TF 版本:2.0.0a
编辑:2.0.0-beta1 也有同样的问题

我保存 pb 的步骤:

  1. 我先设置 K.set_learning_phase(0)
  2. 然后我加载模型 tf.keras.models.load_model
  3. 然后,我定义freeze_session()函数。
  4. (可选我编译模型)
  5. 然后使用该freeze_session()函数tf.keras.backend.get_session

我得到的错误,编译和不编译:

AttributeError:模块“tensorflow.python.keras.api._v2.keras.backend”没有属性“get_session”

我的问题:

  1. TF2没有get_session了?(我知道tf.contrib.saved_model.save_keras_model它不再存在,我也试过tf.saved_model.save哪些没有真正起作用)

  2. 或者get_session只有在我实际训练模型时才有效,而只是加载 h5 不起作用 编辑:同样在新训练的会话中,没有 get_session 可用。

    • 如果是这样,我将如何在未经培训的情况下将 h5 转换为 pb?有好的教程吗?

感谢您的帮助


更新

自 TF2.x 正式发布以来,图形/会话概念发生了变化。该savedmodelAPI应该被使用。您可以将tf.compat.v1.disable_eager_execution()与 TF2.x 一起使用,它将生成一个 pb 文件。但是,我不确定它是哪种 pb 文件类型,因为保存的模型组合从 TF1 更改为 TF2。我会继续挖掘。

小智 6

我做模型保存到pbh5模式:

import logging
import tensorflow as tf
from tensorflow.compat.v1 import graph_util
from tensorflow.python.keras import backend as K
from tensorflow import keras

# necessary !!!
tf.compat.v1.disable_eager_execution()

h5_path = '/path/to/model.h5'
model = keras.models.load_model(h5_path)
model.summary()
# save pb
with K.get_session() as sess:
    output_names = [out.op.name for out in model.outputs]
    input_graph_def = sess.graph.as_graph_def()
    for node in input_graph_def.node:
        node.device = ""
    graph = graph_util.remove_training_nodes(input_graph_def)
    graph_frozen = graph_util.convert_variables_to_constants(sess, graph, output_names)
    tf.io.write_graph(graph_frozen, '/path/to/pb/model.pb', as_text=False)
logging.info("save pb successfully?")
Run Code Online (Sandbox Code Playgroud)

我使用 TF2 来转换模型,如:

  1. 传递keras.callbacks.ModelCheckpoint(save_weights_only=True)model.fit并保存checkpoint,同时培训;
  2. 训练后,self.model.load_weights(self.checkpoint_path)加载checkpoint
  3. self.model.save(h5_path, overwrite=True, include_optimizer=False)另存为h5
  4. 转换h5pb就像上面一样;