def save_mlir(self):
manager = tf.train.CheckpointManager(self.checkpoint,
directory=self.check_dir)
self.checkpoint.restore(manager.latest_checkpoint)
fff = tf.function(self.model).get_concrete_function(tf.TensorSpec([None, 3920]), tf.float32)
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(fff)
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(['pruning', 'function', 'constfold', 'shape', 'remap',
'memory', 'common_subgraph_elimination', 'arithmetic',
'loop', 'dependency', 'debug_stripper']),
graph=frozen_func.graph)
tf_mlir_graph = tf.mlir.experimental.convert_graph_def(graph_def)
outfile = open("./result/kws.mlir", 'wb')
outfile.write(tf_mlir_graph.encode())
outfile.close()
def convert_saved_model_to_graph_def(saved_model_dir_path, model_output_dir_path, signature_name):
imported = tf.saved_model.load(saved_model_dir_path)
f = imported.signatures[signature_name]
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(f, lower_control_flow=False)
input_tensors = [tensor for tensor in frozen_func.inputs if tensor.dtype != tf.resource]
output_tensors = frozen_func.outputs
input_tensor_names = [tensor.name for tensor in frozen_func.inputs if tensor.dtype != tf.resource]
output_tensor_names = [output.name for output in frozen_func.outputs]
print('input_tensor_names:', input_tensor_names)
print('output_tensor_names:', output_tensor_names)
graph_def = run_graph_optimizations(
graph_def,
input_tensors,
output_tensors,
config=get_grappler_config(["constfold", "function"]),
graph=frozen_func.graph)
tf.io.write_graph(graph_or_graph_def=graph_def,
logdir=model_output_dir_path,
name='model_v2.pb',
as_text=False)
print(f'Output model_v2.pb to {model_output_dir_path}/model_v2.pb')
def convert_keras_model(model):
"""Converts a Keras model to TFLite flatbuffer.
Returns:
The converted data in serialized format.
"""
if not tf.executing_eagerly():
raise RuntimeError(
"Graph mode is not supported. Please enable eager execution using "
"tf.enable_eager_execution() when using TensorFlow 1.x")
func = concrete_function_from_keras_model(model)
if version.parse(tf.__version__) >= version.parse("1.15"):
frozen_func = convert_variables_to_constants_v2(
func, lower_control_flow=False)
else:
frozen_func = convert_variables_to_constants_v2(func)
input_tensors = [
tensor for tensor in frozen_func.inputs
if tensor.dtype != tf.dtypes.resource
]
output_tensors = frozen_func.outputs
graph_def = frozen_func.graph.as_graph_def()
# Run a constant folding using grappler since we currently don't implement
# folding for LCE custom ops
graph_def = run_graph_optimizations(
graph_def,
input_tensors,
output_tensors,
config=get_grappler_config(["constfold"]),
graph=frozen_func.graph,
)
# Checks dimensions in input tensor.
for tensor in input_tensors:
# Note that shape_list might be empty for scalar shapes.
shape_list = tensor.shape.as_list()
if None in shape_list[1:]:
raise ValueError(
"None is only supported in the 1st dimension. Tensor '{0}' has "
"invalid shape '{1}'.".format(get_tensor_name(tensor),
shape_list))
elif shape_list and shape_list[0] is None:
# Set the batch size to 1 if undefined.
shape = tensor.shape.as_list()
shape[0] = 1
tensor.set_shape(shape)
return convert_graphdef_to_tflite_flatbuffer(
graph_def.SerializeToString(),
[get_tensor_name(tensor) for tensor in input_tensors],
[
DataType.Name(tensor.dtype.as_datatype_enum)
for tensor in input_tensors
],
[tensor.shape.as_list() for tensor in input_tensors],
[get_tensor_name(tensor) for tensor in output_tensors],
)
def main(_):
if FLAGS.quantize:
try:
_ = tf.contrib
except AttributeError as e:
msg = e.args[0]
msg += ('\n\n The --quantize option still requires contrib, which is not '
'part of TensorFlow 2.0. Please install a previous version:'
'\n `pip install tensorflow<=1.15`')
e.args = (msg,)
raise e
# Create the model and load its weights.
sess = tf.compat.v1.InteractiveSession()
input_tensor, output_tensor = create_inference_graph(
FLAGS.wanted_words, FLAGS.sample_rate, FLAGS.clip_duration_ms,
FLAGS.clip_stride_ms, FLAGS.window_size_ms, FLAGS.window_stride_ms,
FLAGS.feature_bin_count, FLAGS.model_architecture, FLAGS.preprocess)
if FLAGS.quantize:
tf.contrib.quantize.create_eval_graph()
models.load_variables_from_checkpoint(sess, FLAGS.start_checkpoint)
# Turn all the variables into inline constants inside the graph and save it.
frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
sess, sess.graph_def, ['labels_softmax'])
# Here to convert mlir
fff = tf.function(models).get_concrete_function(tf.TensorSpec([-1, models["fingerprint_size"]], tf.float32))
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(fff)
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(
['pruning', 'function', 'constfold', 'shape', 'remap', 'memory', 'common_subgraph_elimination',
'arithmetic', 'loop', 'dependency', 'debug_stripper']),
graph=frozen_func.graph)
tf_mlir_graph = tf.mlir.experimental.convert_graph_def(graph_def)
outfile = open('pix2pix.mlir', 'wb')
outfile.write(tf_mlir_graph.encode())
outfile.close()
def main(_):
"""Export model to MLIR."""
config = get_config(FLAGS.model_name, FLAGS.dataset_cfg, FLAGS.hparam_str)
model = effnetv2_model.EffNetV2Model(FLAGS.model_name, config.model)
# Use call (not build) to match the namescope: tensorflow issues/29576
model(tf.ones([1, 224, 224, 3]), False)
if FLAGS.model_dir:
ckpt = FLAGS.model_dir
if tf.io.gfile.isdir(ckpt):
ckpt = tf.train.latest_checkpoint(FLAGS.model_dir)
utils.restore_tf2_ckpt(model,
ckpt,
exclude_layers=('_head', 'optimizer'))
model.summary()
from tensorflow.lite.python.util import run_graph_optimizations, get_grappler_config
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph
fff = tf.function(model).get_concrete_function(
tf.TensorSpec([1, 224, 224, 3], tf.float32))
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(fff)
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([
'pruning', 'function', 'constfold',
'shape', 'remap', 'memory',
'common_subgraph_elimination',
'arithmetic', 'loop', 'dependency',
'debug_stripper'
]),
graph=frozen_func.graph)
tf_mlir_graph = tf.mlir.experimental.convert_graph_def(graph_def)
print('export model to {}.mlir'.format(FLAGS.model_name))
export_dir = FLAGS.export_dir
if export_dir is None:
export_dir = '.'
os.makedirs(export_dir, exist_ok=True)
outfile = open('{}/{}.mlir'.format(export_dir, FLAGS.model_name), 'wb')
outfile.write(tf_mlir_graph.encode())
outfile.close()
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 to_proto(self):
self.assert_to_proto_paths()
utils.dump2json(obj=vars(self.args),
fp=os.path.join(self.checkpoints_config,
'to_proto_cmd.json'))
model = self.create_model(trainable=False)
self.build_model(model, input_shape=self.input_shape)
model.load_weights(self.checkpoints_restore)
print("[*] Successfully loaded weights!")
# Configure model input
input_dtype = tf.float32
compiled_model = model.get_keras_model(
tf.keras.Input(shape=self.input_shape[1:],
batch_size=1,
name="low"))
compiled_model.summary()
# Freezing model weights, obtaining graph
infer_func = tf.function(compiled_model)
infer_func = infer_func.get_concrete_function(
tf.TensorSpec(shape=compiled_model.inputs[0].shape,
dtype=input_dtype))
frozen_func, graph_def = convert_to_constants(infer_func)
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)
tf.io.write_graph(graph_or_graph_def=graph_def,
logdir=self.checkpoints_to_proto,
name="model.pb",
as_text=False)
def save_frozen_graph(self, path):
real_model_function = tf.function(self.model)
real_model = real_model_function.get_concrete_function(
tf.TensorSpec(self.model.inputs[0].shape,
self.model.inputs[0].dtype))
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
real_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)
tf.io.write_graph(graph_def, './frozen_graph', path)
def _create_frozen_graph(model):
real_model = tf.function(model).get_concrete_function(
tf.TensorSpec(model.inputs[0].shape,
model.inputs[0].dtype,
name=model.inputs[0].name))
frozen_func = convert_variables_to_constants_v2(real_model)
graph_def = frozen_func.graph.as_graph_def()
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 save_mlir(checkpoint, model_func, out_file):
fff = tf.function(model_func).get_concrete_function(tf.TensorSpec([None, 3920]), tf.float32)
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(fff)
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(['pruning', 'function', 'constfold', 'shape', 'remap',
'memory', 'common_subgraph_elimination', 'arithmetic',
'loop', 'dependency', 'debug_stripper']),
graph=frozen_func.graph)
tf_mlir_graph = tf.mlir.experimental.convert_graph_def(graph_def)
outfile = open(out_file, 'wb')
outfile.write(tf_mlir_graph.encode())
outfile.close()