def get_flops(model, config):
"""
Calculate FLOPS for tf.keras.Model or tf.keras.Sequential .
Ignore operations used in only training mode such as Initialization.
Use tf.profiler of tensorflow v1 api.
"""
if not isinstance(model, (tf.keras.Sequential, tf.keras.Model)):
raise KeyError(
"model arguments must be tf.keras.Model or tf.keras.Sequential instance"
)
# convert tf.keras model into frozen graph to count FLOPS about operations used at inference
# FLOPS depends on batch size
inputs = tf.TensorSpec(
[config.batch_size, config.image_size, config.image_size, 3],
tf.float32)
real_model = tf.function(model).get_concrete_function(inputs)
frozen_func, _ = convert_variables_to_constants_v2_as_graph(real_model)
# Calculate FLOPS with tf.profiler
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
cmd="scope",
options=opts)
# print(frozen_func.graph.get_operations())
# TODO: show each FLOPS
return flops.total_float_ops
def freeze(self, func):
"""Freeze the graph."""
# pylint: disable=g-import-not-at-top,disable=g-direct-tensorflow-import
from tensorflow.python.framework.convert_to_constants \
import convert_variables_to_constants_v2_as_graph
_, graphdef = convert_variables_to_constants_v2_as_graph(func)
return graphdef
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 get_flops(model: Union[Model, Sequential],
batch_size: Optional[int] = None) -> int:
"""
Calculate FLOPS for tf.keras.Model or tf.keras.Sequential .
Ignore operations used in only training mode such as Initialization.
Use tf.profiler of tensorflow v1 api.
"""
if not isinstance(model, (Sequential, Model)):
raise KeyError(
"model arguments must be tf.keras.Model or tf.keras.Sequential instanse"
)
if batch_size is None:
batch_size = 1
# convert tf.keras model into frozen graph to count FLOPS about operations used at inference
# FLOPS depends on batch size
inputs = [
tf.TensorSpec([batch_size] + inp.shape[1:], inp.dtype)
for inp in model.inputs
]
real_model = tf.function(model).get_concrete_function(inputs)
frozen_func, _ = convert_variables_to_constants_v2_as_graph(real_model)
# Calculate FLOPS with tf.profiler
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
cmd="scope",
options=opts)
# print(frozen_func.graph.get_operations())
# TODO: show each FLOPS
return flops.total_float_ops
def get_flops(model):
"""
Calculate Flops of model with Tensorflow 1 Profiler
Args:
model: Loaded SavedModel for inference
Returns:
total_flops: Integer with Number of total FLOPs of model
"""
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
tf.TensorSpec(shape=[1, None, None, 3],
dtype=tf.uint8,
name='input_tensor'))
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
full_model)
with tf.Graph().as_default() as graph:
tf.graph_util.import_graph_def(graph_def, name='')
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd="op",
options=opts)
print("Flops: {}".format(flops.total_float_ops))
total_flops = flops.total_float_ops
return total_flops
def get_flops(model, batch_size):
if not isinstance(model, (Sequential, Model)):
raise KeyError(
"model arguments must be tf.keras.Model or tf.keras.Sequential instanse"
)
if batch_size is None:
batch_size = 1
inputs = [
tf.TensorSpec([batch_size] + inp.shape[1:], inp.dtype)
for inp in model.inputs
]
real_model = tf.function(model).get_concrete_function(inputs)
frozen_func, _ = convert_variables_to_constants_v2_as_graph(real_model)
# Calculate FLOPS with tf.profiler
run_meta = tf.compat.v1.RunMetadata()
opts = (tf.compat.v1.profiler.ProfileOptionBuilder().with_empty_output().
build())
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
cmd="scope",
options=opts)
return flops.total_float_ops
def get_flops(model, batch_size=1):
"""
Args:
model: the model to compute FLOPs from
batch_size: batch size (fixed value of 1)
Returns:
Tensorflow profile
"""
if batch_size is None:
batch_size = 1
real_model = tf.function(model).get_concrete_function(
tf.TensorSpec([batch_size] + model.inputs[0].shape[1:],
model.inputs[0].dtype))
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
real_model)
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
cmd='op',
options=opts)
return flops.total_float_ops
def count_flops(model, dummy_inputs, batch_size=None):
"""
Calculate FLOPS for tf.keras.Model or tf.keras.Sequential .
Ignore operations used in only training mode such as Initialization.
Use tf.profiler of tensorflow v1 api.
"""
if not isinstance(model, (keras.Sequential, keras.Model)):
raise ValueError(
"model argument must be tf.keras.Model or tf.keras.Sequential instance"
)
if batch_size is None:
batch_size = 1
if not isinstance(dummy_inputs, list):
dummy_inputs = [dummy_inputs]
# convert tf.keras model into frozen graph to count FLOPS about operations used at inference
# FLOPS depends on batch size
inputs = [
tf.TensorSpec([batch_size] + list(inp.shape[1:]), inp.dtype)
for inp in dummy_inputs
]
real_model = tf.function(model).get_concrete_function(inputs)
frozen_func, _ = convert_variables_to_constants_v2_as_graph(real_model)
# Calculate FLOPS with tf.profiler
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
cmd="scope",
options=opts)
return flops.total_float_ops
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 keras2pb(keras_model,model_path,pb_name):
from tensorflow.python.keras.saving import saving_utils as _saving_utils
from tensorflow.python.framework import convert_to_constants as _convert_to_constants
from tensorflow.compat.v1 import graph_util
func=_saving_utils.trace_model_call(keras_model)
concrete_func=func.get_concrete_function()
frozen_func, graph_def=(_convert_to_constants.convert_variables_to_constants_v2_as_graph(concrete_func,lower_control_flow=False))
graph=graph_util.remove_training_nodes(graph_def)
tf.io.write_graph(graph,model_path,pb_name,as_text=False)
return
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 try_count_flops(model: Union[tf.Module, tf.keras.Model],
inputs_kwargs: Optional[Dict[str, Any]] = None,
output_path: Optional[str] = None):
"""Counts and returns model FLOPs.
Args:
model: A model instance.
inputs_kwargs: An optional dictionary of argument pairs specifying inputs'
shape specifications to getting corresponding concrete function.
output_path: A file path to write the profiling results to.
Returns:
The model's FLOPs.
"""
if hasattr(model, 'inputs'):
try:
# Get input shape and set batch size to 1.
if model.inputs:
inputs = [
tf.TensorSpec([1] + input.shape[1:], input.dtype)
for input in model.inputs
]
concrete_func = tf.function(model).get_concrete_function(
inputs)
# If model.inputs is invalid, try to use the input to get concrete
# function for model.call (subclass model).
else:
concrete_func = tf.function(
model.call).get_concrete_function(**inputs_kwargs)
frozen_func, _ = convert_variables_to_constants_v2_as_graph(
concrete_func)
# Calculate FLOPs.
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
if output_path is not None:
opts['output'] = f'file:outfile={output_path}'
else:
opts['output'] = 'none'
flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph,
run_meta=run_meta,
options=opts)
return flops.total_float_ops
except Exception as e: # pylint: disable=broad-except
logging.info(
'Failed to count model FLOPs with error %s, because the build() '
'methods in keras layers were not called. This is probably because '
'the model was not feed any input, e.g., the max train step already '
'reached before this run.', e)
return None
return None
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 get_flops(model):
concrete = tf.function(lambda inputs: model(inputs))
concrete_func = concrete.get_concrete_function(
[tf.TensorSpec([1, *inputs.shape[1:]]) for inputs in model.inputs])
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
concrete_func)
with tf.Graph().as_default() as graph:
tf.graph_util.import_graph_def(graph_def, name='')
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd="op",
options=opts)
return flops.total_float_ops
def _get_flops(model, input_specs):
"""Credit goes to https://github.com/tensorflow/tensorflow/issues/32809#issuecomment-768977280"""
concrete = tf.function(lambda inputs: model(inputs))
concrete_func = concrete.get_concrete_function(input_specs)
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
concrete_func)
with tf.Graph().as_default() as graph:
tf.graph_util.import_graph_def(graph_def, name="")
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd="op",
options=opts)
return flops.total_float_ops
def constantize(fname):
model = models.load_model(fname)
input_signature = None
# If the model is not a function then the model may include
# a specific batch size, so we include it as well.
if not isinstance(model.call, def_function.Function):
input_signature = saving_utils.model_input_signature(
model, keep_original_batch_size=True)
func = saving_utils.trace_model_call(model, input_signature)
concrete_func = func.get_concrete_function()
_, graph_def = convert.convert_variables_to_constants_v2_as_graph(
concrete_func, lower_control_flow=False)
return graph_def
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 get_flops(model, write_path=tempfile.NamedTemporaryFile().name):
concrete = tf.function(lambda inputs: model(inputs))
concrete_func = concrete.get_concrete_function(
[tf.TensorSpec([1, *inputs.shape[1:]]) for inputs in model.inputs])
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
concrete_func)
with tf.Graph().as_default() as graph:
tf.graph_util.import_graph_def(graph_def, name='')
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
if write_path:
opts['output'] = 'file:outfile={}'.format(
write_path) # suppress output
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd="op",
options=opts)
return flops.total_float_ops
def for_tf2(model):
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))
# for l in full_model.layers:
# print(l)
# real_model.summary()
# Get frozen ConcreteFunction
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
full_model)
# print(graph_def)
print('input: ', graph_def.node[0].op)
print('output: ', graph_def.node[-1].op)
for node in graph_def.node:
# print('#'*50)
# print('name: ' ,node.name)
# print(node.op)
# print([f for f in node.attr.keys()])
if node.op == 'FusedBatchNormV3':
del node.attr["exponential_avg_factor"]
# print([f for f in node.attr.keys()])
# node.attr["exponential_avg_factor"].CopyFrom(
# attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto([1], dtypes.float32)))
# print(frozen_func.graph._collections)
# for f in frozen_func.graph.node:
# print(f)
# 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_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 __get_flops(self, conc_func):
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(
conc_func)
with tf.Graph().as_default() as graph:
tf.compat.v1.graph_util.import_graph_def(graph_def, name='')
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd="op",
options=opts)
flop_tot = flops.total_float_ops
ftxt = open("flops.txt", "w")
for idx, name in enumerate(['', 'K', 'M', 'G', 'T']):
text = '%.3f [%sFLOPS]' % (flop_tot / 10**(3 * idx), name)
print(text)
ftxt.write("%s\n" % (text))
ftxt.close()
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()
def main(_):
logging.set_verbosity(logging.DEBUG if FLAGS.debug else logging.INFO)
params = Config(FLAGS.config_path).params
trainer = Trainer(
run_mode='export', # noqa: F841
strategy=tf.distribute.OneDeviceStrategy(
device='/cpu:0'), # noqa: E501
model_fn=model_builder(params),
train_input_fn=None,
val_input_fn=None,
style_loss_weight=params.training.style_loss_weight,
content_loss_weight=params.training.content_loss_weight,
train_steps=params.training.train_steps,
val_steps=params.training.validation_steps,
val_freq=params.training.validation_freq,
steps_per_execution=params.training.steps_per_execution,
batch_size=params.training.batch_size,
model_dir=params.experiment.model_dir,
save_every=params.training.save_every,
restore_checkpoint=params.training.restore_checkpoint,
summary_dir=params.experiment.tensorboard_dir,
name=params.experiment.name)
inference_model = trainer.model
inference_model.build([(1, None, None, 3), (1, None, None, 3), ()])
inference_model.optimizer = None
inference_model.compiled_loss = None
inference_model.compiled_metrics = None
inference_model._metrics = []
input_signature = {
'style_images':
tf.TensorSpec(shape=[1, None, None, 3],
name='style_images',
dtype=tf.float32),
'content_images':
tf.TensorSpec(shape=[1, None, None, 3],
name='content_images',
dtype=tf.float32),
'alpha':
tf.TensorSpec(shape=[], name='alpha', dtype=tf.float32),
'resize':
tf.TensorSpec(shape=[], name='resize', dtype=tf.bool)
}
logging.debug(
'Signature for serving_default :\n{}'.format(input_signature))
inference_model._saved_model_inputs_spec = input_signature
preprocessing_pipeline = PreprocessingPipeline(params=params,
is_validation_dataset=True)
@tf.function
def serving_fn(sample):
style_images = sample['style_images']
content_images = sample['content_images']
if sample['resize']:
style_images = preprocessing_pipeline.inference_pipeline(
images=style_images)
content_images = preprocessing_pipeline.inference_pipeline(
images=content_images)
stylized_images = inference_model.call(
(style_images, content_images, sample['alpha']),
training=False)['synthesized_images']
stylized_images = preprocessing_pipeline.denormalize(stylized_images)
return {'stylized_images': stylized_images}
frozen_serving_fn, _ = convert_variables_to_constants_v2_as_graph(
serving_fn.get_concrete_function(input_signature),
aggressive_inlining=True)
class InferenceModule(tf.Module):
def __init__(self, inference_function):
super(InferenceModule, self).__init__(name='inference_module')
self.inference_function = inference_function
@tf.function
def run_inference(self, sample):
outputs = self.inference_function(**sample)
return {'stylized_images': outputs[0]}
inference_module = InferenceModule(inference_function=frozen_serving_fn)
signatures = {
'serving_default':
inference_module.run_inference.get_concrete_function(input_signature)
}
print(signatures['serving_default'].output_shapes)
for _signature_name, _concrete_fn in signatures.items():
input_shapes = {
x.name.split(':')[0]: x.shape.as_list()
for x in _concrete_fn.inputs
}
output_shapes = {
k: v.as_list()
for k, v in _concrete_fn.output_shapes.items()
}
logging.info(
'\nSignature: {}\n Input Shapes:\n {}\nOutput Shapes:\n{}'.format(
_signature_name, input_shapes, output_shapes))
logging.info('Exporting `saved_model` to {}'.format(FLAGS.export_dir))
tf.saved_model.save(obj=inference_module,
export_dir=FLAGS.export_dir,
signatures=signatures,
options=tf.saved_model.SaveOptions(
experimental_custom_gradients=False))
parser.add_argument('--filters',
type=int,
default=128,
help='Number of filters')
parser.add_argument('--backbone',
type=str,
default=list(_BACKBONES)[0],
choices=list(_BACKBONES),
help='Backbone network')
args = parser.parse_args()
model = create_ksac_net(args.input_shape + [3],
args.n_classes,
args.filters,
backbone=args.backbone)
concrete = tf.function(lambda inputs: model(inputs))
concrete_func = concrete.get_concrete_function(
[tf.TensorSpec([1] + args.input_shape + [3])])
_, graph_def = convert_variables_to_constants_v2_as_graph(concrete_func)
with tf.Graph().as_default() as graph:
tf.graph_util.import_graph_def(graph_def, name='')
run_meta = tf.compat.v1.RunMetadata()
opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
flops = tf.compat.v1.profiler.profile(graph=graph,
run_meta=run_meta,
cmd='op',
options=opts)
print(model.summary())
print(f'{flops.total_float_ops:,} FLOPs')
tf.keras.utils.plot_model(model, show_shapes=True)
def convert(self):
"""Converts a TensorFlow GraphDef based on instance variables.
Returns:
The converted data in serialized format.
Raises:
ValueError:
Multiple concrete functions are specified.
Input shape is not specified.
Invalid quantization parameters.
"""
# TODO(b/130297984): Add support for converting multiple function.
if len(self._funcs) != 1:
raise ValueError("This converter can only convert a single "
"ConcreteFunction. Converting multiple functions is "
"under development.")
# graph_def is used here to preserve the node bug information
frozen_func, graph_def = (
_convert_to_constants.convert_variables_to_constants_v2_as_graph(
self._funcs[0], lower_control_flow=False))
input_tensors = [
tensor for tensor in frozen_func.inputs
if tensor.dtype != _dtypes.resource
]
output_tensors = frozen_func.outputs
# Run a Grappler pass.
graph_def = _run_graph_optimizations(
graph_def,
input_tensors,
output_tensors,
config=self._grappler_config(),
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)
self._validate_quantization()
self._validate_representative_dataset()
if self._trackable_obj is None:
self._debug_info = _get_debug_info(
_build_debug_info_func(self._funcs[0].graph), graph_def)
else:
self._debug_info = _get_debug_info(
_convert_debug_info_func(self._trackable_obj.graph_debug_info),
graph_def)
converter_kwargs = self._get_base_converter_args()
if not self.experimental_new_converter:
logging.warning(
"Please consider switching to use new converter by setting "
"experimental_new_converter to true. "
"Old converter (TOCO) is deprecated and flow will be switched on "
"by default to use new converter soon.")
else:
logging.info("Using experimental converter: If you encountered a problem "
"please file a bug. You can opt-out "
"by setting experimental_new_converter=False")
# Converts model.
result = _toco_convert_impl(
input_data=graph_def,
input_tensors=input_tensors,
output_tensors=output_tensors,
**converter_kwargs)
if self._is_calibration_quantize():
result = self._calibrate_quantize_model(
result, constants.FLOAT, constants.FLOAT,
self.experimental_new_quantizer)
return result
