def tf1_export_ema_ckpt():
"""Restore variables from a given checkpoint."""
with tf1.Session() as sess:
model = effnetv2_model.EffNetV2Model(FLAGS.model_name,
FLAGS.hparam_str)
batch_size = FLAGS.batch_size
isize = FLAGS.image_size or model.cfg.eval.isize
inputs = tf.ones((batch_size, isize, isize, 3), tf.float32)
_ = model(inputs, training=False)
sess.run(tf1.global_variables_initializer())
if tf.io.gfile.isdir(FLAGS.model_dir):
ckpt_path = tf1.train.latest_checkpoint(FLAGS.model_dir)
else:
ckpt_path = FLAGS.model_dir
ema = tf1.train.ExponentialMovingAverage(decay=0.0)
ema_vars = utils.get_ema_vars()
var_dict = ema.variables_to_restore(ema_vars)
ema_assign_op = ema.apply(ema_vars)
tf1.train.get_or_create_global_step()
sess.run(tf1.global_variables_initializer())
saver = tf1.train.Saver(var_dict, max_to_keep=1)
# Restore all variables from ckpt.
saver.restore(sess, ckpt_path)
print('export model to {}'.format(FLAGS.export_dir))
sess.run(ema_assign_op)
saver = tf1.train.Saver(max_to_keep=1, save_relative_paths=True)
saver.save(sess, FLAGS.export_dir)
def tf1_benchmark():
"""Run TF1 inference and benchmark."""
# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
from tensorflow.python.client import timeline
with tf1.Session() as sess:
model = effnetv2_model.EffNetV2Model(FLAGS.model_name,
FLAGS.hparam_str)
batch_size = FLAGS.batch_size
run_options = tf1.RunOptions(trace_level=tf1.RunOptions.FULL_TRACE)
run_metadata = tf1.RunMetadata()
isize = FLAGS.image_size or model.cfg.eval.isize
data_dtype = tf.float16 if FLAGS.mixed_precision else tf.float32
inputs = tf.ones((batch_size, isize, isize, 3), data_dtype)
output = model(inputs, training=False)
sess.run(tf1.global_variables_initializer())
print('starting warmup.')
for _ in range(5):
sess.run(output)
print('starting benchmark.')
start = time.perf_counter()
for _ in range(10):
sess.run(output)
end = time.perf_counter()
inference_time = (end - start) / 10
print('Per batch inference time: ', inference_time)
print('FPS: ', batch_size / inference_time)
if FLAGS.trace_file:
sess.run(output, options=run_options, run_metadata=run_metadata)
with tf.io.gfile.GFile(FLAGS.trace_file, 'w') as f:
trace = timeline.Timeline(step_stats=run_metadata.step_stats)
f.write(trace.generate_chrome_trace_format(show_memory=True))
def build_tf2_model():
"""Build the tf2 model."""
tf.config.run_functions_eagerly(FLAGS.debug)
config = get_config(FLAGS.model_name, FLAGS.dataset_cfg, FLAGS.hparam_str)
if config.runtime.mixed_precision:
# Use 'mixed_float16' if running on GPUs.
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
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)
model.load_weights(ckpt)
model.summary()
class ExportModel(tf.Module):
"""Export a saved model."""
def __init__(self, model):
super().__init__()
self.model = model
@tf.function
def f(self, images):
return self.model(images, training=False)[0]
return ExportModel(model)
def build_model(in_images):
"""Build model using the model_name given through the command line."""
config.model.num_classes = config.data.num_classes
model = effnetv2_model.EffNetV2Model(config.model.model_name,
config.model)
logits = model(in_images, training=is_training)[0]
return logits
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 test_effnetv2(self, model_name, expected_params):
images = tf.zeros((10, 224, 224, 3), dtype=tf.float32)
model = effnetv2_model.EffNetV2Model(model_name)
_ = model(images)
self.assertEqual(model.count_params(), expected_params)