Convert frozen model(.pb) to savedmodel - tensorflow

Question

Recently I tried to convert the model (tf1.x) to the saved_model, and followed the official migrate document. However in my use case, most of model in my hand or tensorflow model zoo usually is pb file, and according to the official document says that
There is no straightforward way to upgrade a raw Graph.pb file to TensorFlow 2.0, but if you have a "Frozen graph" (a tf.Graph where the variables have been turned into constants), then it is possible to convert this to a concrete_function using v1.wrap_function:
But I still do not understand how to converted to saved_model format.

Answer 1

in TF1 mode:
import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
def convert_pb_to_server_model(pb_model_path, export_dir, input_name='input:0', output_name='output:0'):
graph_def = read_pb_model(pb_model_path)
convert_pb_saved_model(graph_def, export_dir, input_name, output_name)
def read_pb_model(pb_model_path):
with tf.gfile.GFile(pb_model_path, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
def convert_pb_saved_model(graph_def, export_dir, input_name='input:0', output_name='output:0'):
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
sigs = {}
with tf.Session(graph=tf.Graph()) as sess:
tf.import_graph_def(graph_def, name="")
g = tf.get_default_graph()
inp = g.get_tensor_by_name(input_name)
out = g.get_tensor_by_name(output_name)
sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
tf.saved_model.signature_def_utils.predict_signature_def(
{"input": inp}, {"output": out})
builder.add_meta_graph_and_variables(sess,
[tag_constants.SERVING],
signature_def_map=sigs)
builder.save()
in TF2 mode:
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph
from tensorflow.lite.python.util import run_graph_optimizations, get_grappler_config
import numpy as np
def frozen_keras_graph(func_model):
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(func_model)
input_tensors = [
tensor for tensor in frozen_func.inputs
if tensor.dtype != tf.resource
]
output_tensors = frozen_func.outputs
graph_def = run_graph_optimizations(
graph_def,
input_tensors,
output_tensors,
config=get_grappler_config(["constfold", "function"]),
graph=frozen_func.graph)
return graph_def
def convert_keras_model_to_pb():
keras_model = train_model()
func_model = tf.function(keras_model).get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
graph_def = frozen_keras_graph(func_model)
tf.io.write_graph(graph_def, '/tmp/tf_model3', 'frozen_graph.pb')
def convert_saved_model_to_pb():
model_dir = '/tmp/saved_model'
model = tf.saved_model.load(model_dir)
func_model = model.signatures["serving_default"]
graph_def = frozen_keras_graph(func_model)
tf.io.write_graph(graph_def, '/tmp/tf_model3', 'frozen_graph.pb')
Or:
def convert_saved_model_to_pb(output_node_names, input_saved_model_dir, output_graph_dir):
from tensorflow.python.tools import freeze_graph
output_node_names = ','.join(output_node_names)
freeze_graph.freeze_graph(input_graph=None, input_saver=None,
input_binary=None,
input_checkpoint=None,
output_node_names=output_node_names,
restore_op_name=None,
filename_tensor_name=None,
output_graph=output_graph_dir,
clear_devices=None,
initializer_nodes=None,
input_saved_model_dir=input_saved_model_dir)
def save_output_tensor_to_pb():
output_names = ['StatefulPartitionedCall']
save_pb_model_path = '/tmp/pb_model/freeze_graph.pb'
model_dir = '/tmp/saved_model'
convert_saved_model_to_pb(output_names, model_dir, save_pb_model_path)

Answer 2

To ensure whether my understanding is correct, so I also post what I learnt:
If anyone want to migrate tf1.x to tf2.x, please followed the official post first.
In tensorflow 2.0, tf.train.Saver and freeze_graph was already replaced by saved_model.
And if anyone want to convert pb model from tf1.x into saved_model, you can followed #Boluoyu answer. But if your run-time environment is above tf2.0, you can use following code:
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
def covert_pb_to_server_model(pb_model_path, export_dir, input_name='input', output_name='output'):
graph_def = read_pb_model(pb_model_path)
covert_pb_saved_model(graph_def, export_dir, input_name, output_name)
def read_pb_model(pb_model_path):
with tf.gfile.GFile(pb_model_path, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
def covert_pb_saved_model(graph_def, export_dir, input_name='input', output_name='output'):
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
sigs = {}
with tf.Session(graph=tf.Graph()) as sess:
tf.import_graph_def(graph_def, name="")
g = tf.get_default_graph()
inp = g.get_tensor_by_name(input_name)
out = g.get_tensor_by_name(output_name)
sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
tf.saved_model.signature_def_utils.predict_signature_def(
{"input": inp}, {"output": out})
builder.add_meta_graph_and_variables(sess,
[tag_constants.SERVING],
signature_def_map=sigs)
builder.save()