Tensorflow 2.0将keras模型转换为.pb文件

时间:2020-01-31 14:18:32

标签: python python-3.x tensorflow keras

我正在生成一个keras模型并将其保存到.h5文件,然后尝试将其转换为.pb文件以供以后统一使用。

我遵循了convert tensorflow model to pb tensorflow的一些说明以及一些其他建议,这些建议似乎可以追溯到tensorflow 1.0是最新版本时,但它们也存在类似的问题。

我在下面的代码中遇到的错误是当我尝试将变量转换为常量时:它抱怨我的变量不在会话定义的图中。 (我不是张量流的菜鸟,所以我并不完全知道这是什么意思,但我认为这与我的模型无关)。

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
from keras import backend as K

tf.keras.backend.set_learning_phase(0)

pre_model = tf.keras.models.load_model("final_model.h5")

print(pre_model.inputs)
print(pre_model.outputs)

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    """
    Freezes the state of a session into a pruned computation graph.

    Creates a new computation graph where variable nodes are replaced by
    constants taking their current value in the session. The new graph will be
    pruned so subgraphs that are not necessary to compute the requested
    outputs are removed.
    @param session The TensorFlow session to be frozen.
    @param keep_var_names A list of variable names that should not be frozen,
                          or None to freeze all the variables in the graph.
    @param output_names Names of the relevant graph outputs.
    @param clear_devices Remove the device directives from the graph for better portability.
    @return The frozen graph definition.
    """
    from tensorflow.compat.v1.graph_util import convert_variables_to_constants
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.compat.v1.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.compat.v1.global_variables()]
        # Graph -> GraphDef ProtoBuf
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = convert_variables_to_constants(session, input_graph_def,
                                                      output_names, freeze_var_names)
        return frozen_graph


frozen_graph = freeze_session(tf.compat.v1.keras.backend.get_session(), output_names=[out.op.name for out in pre_model.outputs])

有输出+错误:

[<tf.Tensor 'conv2d_1_input:0' shape=(None, 28, 28, 1) dtype=float32>]
[<tf.Tensor 'dense_2/Identity:0' shape=(None, 10) dtype=float32>]
File "saveGraph.py", line 40, in freeze_session
    output_names, freeze_var_names)
  File "C:\Users\jgoer\AppData\Roaming\Python\Python37\site-packages\tensorflow_core\python\util\deprecation.py", line 324, in new_func
    return func(*args, **kwargs)
  File "C:\Users\jgoer\AppData\Roaming\Python\Python37\site-packages\tensorflow_core\python\framework\graph_util_impl.py", line 277, in convert_variables_to_constants
    inference_graph = extract_sub_graph(input_graph_def, output_node_names)
  File "C:\Users\jgoer\AppData\Roaming\Python\Python37\site-packages\tensorflow_core\python\util\deprecation.py", line 324, in new_func
    return func(*args, **kwargs)
  File "C:\Users\jgoer\AppData\Roaming\Python\Python37\site-packages\tensorflow_core\python\framework\graph_util_impl.py", line 197, in extract_sub_graph    
    _assert_nodes_are_present(name_to_node, dest_nodes)
File "C:\Users\jgoer\AppData\Roaming\Python\Python37\site-packages\tensorflow_core\python\framework\graph_util_impl.py", line 152, in _assert_nodes_are_present
    assert d in name_to_node, "%s is not in graph" % d
AssertionError: dense_2/Identity is not in graph

1 个答案:

答案 0 :(得分:2)

看看TensorFlow的tutorial on saving and loading models。您可以使用model.save("path"),如果不包括扩展名,则模型将以SavedModel格式保存。

import tensorflow as tf

pre_model = tf.keras.models.load_model("final_model.h5")
pre_model.save("saved_model")