将冷冻的TensorFlow pb包装在tf.keras模型中

时间:2018-09-19 00:08:22

标签: python tensorflow keras

我正在尝试在更大的tf.keras培训管道中使用经过冻结的,经过预训练的DeepLabv3模型,但是一直难以弄清如何将其用作tf.keras模型。我正在尝试使用tf.keras,因为在多个前向传递中间,使用feed_dict(我知道使用冻结图的唯一方法)会降低速度。以下代码中引用的deeplab模型是在常规keras(与tf.contrib.keras相反)中构建的

from keras import backend as K

# Create, compile and train model...

frozen_graph = freeze_session(K.get_session(),
                        output_names=[out.op.name for out in deeplab.outputs])
tf.train.write_graph(frozen_graph, "./", "my_model.pb", as_text=False)
graph = load_graph("my_model.pb")

# We can verify that we can access the list of operations in the graph
for op in graph.get_operations():
    print(op.name)
    # prefix/Placeholder/inputs_placeholder
    # ...
    # prefix/Accuracy/predictions

# We access the input and output nodes 
x = graph.get_tensor_by_name("prefix/input_1:0")
y = graph.get_tensor_by_name("prefix/bilinear_upsampling_2/ResizeBilinear:0")

# We launch a Session
with tf.Session(graph=graph) as sess:
    print(graph)
    model2 = models.Model(inputs=x,outputs=y)
    model2.summary()

我得到一个错误

ValueError: Input tensors to a Model must come from `tf.layers.Input`. Received: Tensor("prefix/input_1:0", shape=(?, 512, 512, 3), dtype=float32) (missing previous layer metadata).

我觉得我已经看到其他人用输入层替换输入张量来欺骗tf.keras来构建图形,但是几个小时后,我感到卡住了。任何帮助将不胜感激!

1 个答案:

答案 0 :(得分:0)

您可以从其config重新创建模型对象。在https://keras.io/models/about-keras-models/处参见from_config方法。

该配置由save_model/load_model functions存储和加载回。我对freeze_session不熟悉。