如何将Keras模型保存为冻结图? [英] How to save Keras model as frozen graph?

问题描述

我正在使用Tensorflow 2.0，并希望将以下Keras模型存储为冻结图.

I am working with Tensorflow 2.0 and want to store the following Keras model as frozen graph.

import tensorflow as tf
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(64, input_shape=[100]))
model.add(tf.keras.layers.Dense(32, activation='relu'))
model.add(tf.keras.layers.Dense(16, activation='relu'))
model.add(tf.keras.layers.Dense(2, activation='softmax'))
model.summary()
model.save('./models/')

我在Tensorflow 2.0中找不到如何执行此操作的好示例.我在其中找到了 freeze_graph.py 文件Tensorflow Github存储库，但是很难把它包裹住.

I can't find any good examples how to do this in Tensorflow 2.0. I have found the freeze_graph.py file in the Tensorflow Github repository but find it hard to wrap my head around it.

我使用以下方法加载上述文件:

I load the file mentioned above using:

from tensorflow.python.tools.freeze_graph import freeze_graph

但是我到底要为freeze_graph函数本身提供什么呢?在这里，我用问号标记了不确定的参数.

But what exactly do I have to provide to the freeze_graph function itself? Here I marked the arguments where I am not sure with a questionmark.

freeze_graph(input_graph=?,
             input_saver='',
             input_binary=False,
             input_checkpoint=?,
             output_node_names=?,
             restore_op_name='',
             filename_tensor_name='',
             output_graph='./frozen_graph.pb',
             clear_devices=True,
             initializer_nodes='')

有人可以提供一个简单的示例来说明如何使用freeeze_graph函数将上面的模型存储为冻结图吗?

Can someone provide a simple example that shows how I can store the model above as a frozen graph using the freeeze_graph function?

推荐答案

Freeze_Graph现在已在Tensorflow 2.0中消失.
您可以在此处进行检查

Freeze_Graph is now gone in Tensorflow 2.0.
You can check it here Tensorflow 2.0 : frozen graph support.

代码中的 .save 方法除外.
.save 方法已经保存了一个.pb以便进行推理. 另外，您也可以使用以下代码.

Except for the .save method that you have in your code.
.save Method is already saving a .pb ready for inference. As an alternative, you can also use the below code.

您还可以使用 convert_variables_to_constants_v2

You can also use convert_variables_to_constants_v2

下面是示例代码.

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

model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(64, input_shape=(1,)))
model.add(tf.keras.layers.Dense(32, activation='relu'))
model.add(tf.keras.layers.Dense(16, activation='relu'))
model.add(tf.keras.layers.Dense(1, activation='softmax'))
model.compile(optimizer='adam', loss='mse')
model.summary()

# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name="yourInputName"))
# 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 to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir="./frozen_models",
                  name="frozen_graph.pb",
                  as_text=False)

### USAGE ##
def wrap_frozen_graph(graph_def, inputs, outputs, print_graph=False):
    def _imports_graph_def():
        tf.compat.v1.import_graph_def(graph_def, name="")

    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
    import_graph = wrapped_import.graph

    print("-" * 50)
    print("Frozen model layers: ")
    layers = [op.name for op in import_graph.get_operations()]
    if print_graph == True:
        for layer in layers:
            print(layer)
    print("-" * 50)

    return wrapped_import.prune(
        tf.nest.map_structure(import_graph.as_graph_element, inputs),
        tf.nest.map_structure(import_graph.as_graph_element, outputs))

## Example Usage ###
# Load frozen graph using TensorFlow 1.x functions
with tf.io.gfile.GFile("./frozen_models/frozen_graph.pb", "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    loaded = graph_def.ParseFromString(f.read())

# Wrap frozen graph to ConcreteFunctions
frozen_func = wrap_frozen_graph(graph_def=graph_def,
                                inputs=["yourInputName:0"],
                                outputs=["Identity:0"],
                                print_graph=True)
print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
# Get predictions for test images
predictions = frozen_func(yourInputName=tf.constant([[3.]]))
# Print the prediction for the first image
print("-" * 50)
print("Example prediction reference:")
print(predictions[0].numpy())

这篇关于如何将Keras模型保存为冻结图?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！