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