def from_function(func, input_names, output_names, large_model=False): if large_model: return convert_variables_to_constants_large_model(func) try: if get_tf_version() < LooseVersion("2.2"): frozen_func = convert_variables_to_constants_v2( func, lower_control_flow=False) else: frozen_func = convert_variables_to_constants_v2( func, lower_control_flow=False, aggressive_inlining=True) except ValueError as e: if "incompatible with expected resource" in str(e): frozen_func = convert_variables_to_constants_large_model(func) logger.warning( "TF freezing failed. Attempting to fix freezing errors.") graph_def = fix_freezing_errors(frozen_func) else: raise e else: graph_def = frozen_func.graph.as_graph_def(add_shapes=True) # output_names = [i.name for i in frozen_func.outputs] with tf.Graph().as_default() as tf_graph: with tf_session(graph=tf_graph) as sess: tf.import_graph_def(graph_def, name='') input_names = inputs_without_resource(sess, input_names) graph_def = tf_optimize(input_names, output_names, graph_def) return graph_def
def _freeze_saved_model_v2(concrete_func, control_flow_v2=False): if tf.__version__ < '2.2.0': return convert_to_constants.convert_variables_to_constants_v2( concrete_func, lower_control_flow=not control_flow_v2).graph return convert_to_constants.convert_variables_to_constants_v2( concrete_func, lower_control_flow=not control_flow_v2, aggressive_inlining=True).graph
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 testConstSavedModel(self): """Test a basic model with functions to make sure functions are inlined.""" input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.f = def_function.function(lambda x: 2. * x) to_save = root.f.get_concrete_function(input_data) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save(root, save_dir, to_save) saved_model = load(save_dir) input_func = saved_model.signatures["serving_default"] variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(variable_graph_def)) self.assertTrue(variable_graph_def.library.function) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(constant_graph_def.library.function) # Check value. expected_value = root.f(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def testKerasModel(self): input_data = constant_op.constant(1., shape=[1, 1]) # Create a simple Keras model. x = [-1, 0, 1, 2, 3, 4] y = [-3, -1, 1, 3, 5, 7] model = keras.models.Sequential( [keras.layers.Dense(units=1, input_shape=[1])]) model.compile(optimizer="sgd", loss="mean_squared_error") model.fit(x, y, epochs=1) # Get the concrete function from the Keras model. @def_function.function def to_save(x): return model(x) input_func = to_save.get_concrete_function(input_data) variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(2, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = to_save(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def testConstructConcreteFunction(self): input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.v1 = variables.Variable(3.) root.v2 = variables.Variable(2.) root.f = def_function.function(lambda x: root.v1 * root.v2 * x) func = root.f.get_concrete_function(input_data) input_func = convert_to_constants._construct_concrete_function( func, func.graph.as_graph_def()) # Test if model has enough metadata to be frozen afterwards. variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(2, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = root.f(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def save_model(path, name, model, suffix='tf'): """ Save model for h5/json/pb/tf/frozen :param path: str. :param name: str. :param model: tf.keras.models.Model. :param suffix: str. For different format. """ if not os.path.exists(path): os.mkdir(path) if suffix == 'h5': model.save(path + "/" + name + ".h5") # tf.keras.models.save_model(model, path, save_format='h5') elif suffix == 'json': model_json = model.to_json() with open(path + "/" + name + ".json", "w") as json_file: json_file.write(model_json) model.save_weights(path + "/" + name + ".h5") elif suffix == 'pb' or suffix == 'tf': tf.keras.models.save_model(model, path, save_format='tf') elif suffix == 'frozen': full_model = tf.function(lambda x: model(x)) full_model = full_model.get_concrete_function( x=tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype)) frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=path, name=name + '.pb', as_text=False) else: raise TypeError('suffix=\'%s\'' % suffix)
def testMultiFunctionModel(self): """Test a basic model with Variables.""" class BasicModel(tracking.AutoTrackable): def __init__(self): self.y = None self.z = None @def_function.function def add(self, x): if self.y is None: self.y = variables.Variable(2.) return x + self.y @def_function.function def sub(self, x): if self.z is None: self.z = variables.Variable(3.) return x - self.z input_data = {"x": constant_op.constant(1., shape=[1])} root = BasicModel() input_func = root.add.get_concrete_function(input_data["x"]) variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(1, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) self._testConvertedFunction(root, root.add, output_func, input_data)
def testVariableSavedModel(self): """Test a basic model with Variables with saving/loading the SavedModel.""" input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.v1 = variables.Variable(3.) root.v2 = variables.Variable(2.) root.f = def_function.function(lambda x: root.v1 * root.v2 * x) to_save = root.f.get_concrete_function(input_data) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save(root, save_dir, to_save) saved_model = load(save_dir) input_func = saved_model.signatures["serving_default"] variable_graph_def = input_func.graph.as_graph_def() self.assertTrue(self._hasStatefulPartitionedCallOp(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = root.f(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def freezeModel(model): # Convert Keras model to ConcreteFunction 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)) # Get frozen ConcreteFunction frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() layers = [op.name for op in frozen_func.graph.get_operations()] print("-" * 50) print("Frozen model layers: ") for layer in layers: print(layer) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Save frozen graph from frozen ConcreteFunction to hard drive tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir="./frozen_models", name="frozen_graph.pb", as_text=False)
def freeze_model(self): """ Freezes the model """ tf.saved_model.save(self.tf_model, self.helpers.confs["model"]["saved_model_dir"]) fmodel = tf.function(lambda x: self.tf_model(x)) fmodel = fmodel.get_concrete_function(x=tf.TensorSpec( self.tf_model.inputs[0].shape, self.tf_model.inputs[0].dtype)) freeze = convert_variables_to_constants_v2(fmodel) freeze.graph.as_graph_def() layers = [op.name for op in freeze.graph.get_operations()] self.helpers.logger.info("Frozen model layers") for layer in layers: self.helpers.logger.info(layer) self.helpers.logger.info("Frozen model inputs") self.helpers.logger.info(freeze.inputs) self.helpers.logger.info("Frozen model outputs") self.helpers.logger.info(freeze.outputs) tf.io.write_graph( graph_or_graph_def=freeze.graph, logdir=self.helpers.confs["model"]["freezing_log_dir"], name=self.helpers.confs["model"]["frozen"], as_text=False)
def get_func_from_saved_model(saved_model_dir): saved_model_loaded = tf.saved_model.load( saved_model_dir, tags=[tag_constants.SERVING]) graph_func = saved_model_loaded.signatures[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] graph_func = convert_to_constants.convert_variables_to_constants_v2(graph_func) return graph_func
def testStatelessIf(self): """Test whether StatelessIf op freezes correctly.""" input_data = {"b": constant_op.constant(True)} x = constant_op.constant([1., 2.], shape=[1, 2], name="x") def true_fn(): return x def false_fn(): return x + 2 @def_function.function(input_signature=[ tensor_spec.TensorSpec(shape=(), dtype=dtypes.bool) ]) def model(b): return cond_v2.cond_v2(b, true_fn, false_fn) root = tracking.AutoTrackable() root.f = model input_func = root.f.get_concrete_function() input_func(**input_data) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func, lower_control_flow=False) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) self._testConvertedFunction(root, root.f, output_func, input_data)
def _graph_def_from_saved_model_or_keras_model(filename): """ Utility function that returns GraphDef object from the given SavedModel or HDF5 model. :param filename: TensorFlow SavedModel directory or Keras HDF5 model (.h5) file. :return: TensorFlow GraphDef object. """ try: import tensorflow as tf from tensorflow.python.keras.saving import saving_utils as _saving_utils from tensorflow.python.framework import convert_to_constants as _convert_to_constants model = tf.keras.models.load_model(filename) tf.keras.backend.set_learning_phase(False) func = _saving_utils.trace_model_call(model) concrete_func = func.get_concrete_function() # concrete_func = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] frozen_func = _convert_to_constants.convert_variables_to_constants_v2( concrete_func) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) except ImportError as e: raise ImportError( 'Failed to import TensorFlow utilities. {}.'.format(e)) except Exception as e: raise RuntimeError( 'Failed to load SavedModel or .h5 model. {}.'.format(e)) return graph_def
def testLoop(self): input_data = { "x": constant_op.constant([1., 2., 3., 4.], shape=[2, 2]) } weights = variables.Variable([[0.1, 0.2], [0.3, 0.4]], dtype=dtypes.float32) def condition(x): return math_ops.reduce_sum(x) < 100 def body(x): return math_ops.add(x, weights) @def_function.function(input_signature=[ tensor_spec.TensorSpec(shape=[2, 2], dtype=dtypes.float32) ]) def model(x): return control_flow_ops.while_loop(condition, body, [x]) root = tracking.AutoTrackable() root.f = model input_func = root.f.get_concrete_function() input_func(**input_data) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func, lower_control_flow=False) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) self._testConvertedFunction(root, root.f, output_func, input_data)
def __init__(self, keras_model_path, inputshape, in_nodes, dest_nodes): if LooseVersion(tensorflow.__version__) < LooseVersion('1.8.0'): raise ImportError( 'Your TensorFlow version %s is outdated. ' 'MMdnn requires tensorflow>=1.8.0' % tensorflow.__version__) super(TensorflowParser2, self).__init__() self.weight_loaded = True import tensorflow as tf from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 model = tf.keras.models.load_model(keras_model_path, compile=False) 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)) frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() g = frozen_func.graph from tensorflow.python.client.session import Session from tensorflow.python.training.saver import export_meta_graph with Session(graph = g) as sess: tempdir = tempfile.mkdtemp() meta_graph_def = export_meta_graph(filename=os.path.join(tempdir, 'my-model.meta')) model = meta_graph_def.graph_def shutil.rmtree((tempdir)) self.tf_graph = TensorflowGraph(model) self.tf_graph.build()
def to_frozen_graph(model: tf.keras.Model): """ Returns a frozen graph def for a tf.keras model. :param model: The tf.keras model we want to convert. :return: The frozen graph def for the provided model. """ func = saving_utils.trace_model_call(model) concrete_func = func.get_concrete_function() graph_captures = concrete_func.graph._captures # pylint: disable=protected-access captured_inputs = [t_name.name for t_val, t_name in graph_captures.values()] input_names = [input_tensor.name for input_tensor in concrete_func.inputs if input_tensor.name not in captured_inputs] output_names = [output_tensor.name for output_tensor in concrete_func.outputs if output_tensor.dtype != tf.dtypes.resource] with tf.device("/cpu:0"): frozen_func = convert_variables_to_constants_v2(concrete_func, lower_control_flow=False, aggressive_inlining=True) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) with tf.Graph().as_default(): # pylint: disable=not-context-manager tf.import_graph_def(graph_def, name='') frozen_graph = tf_optimize_grappler(input_names, output_names, graph_def) return frozen_graph
def _convert_saved_model_v2(self): """Convert the input SavedModel in 2.0 format.""" assert context.executing_eagerly() self._saved_model = load.load(self._input_saved_model_dir, self._input_saved_model_tags) func = self._saved_model.signatures[ self._input_saved_model_signature_key] frozen_func = convert_to_constants.convert_variables_to_constants_v2( func) self._grappler_meta_graph_def = saver.export_meta_graph( graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) # Add a collection 'train_op' so that Grappler knows the outputs. fetch_collection = meta_graph_pb2.CollectionDef() for array in frozen_func.inputs + frozen_func.outputs: fetch_collection.node_list.value.append(array.name) self._grappler_meta_graph_def.collection_def["train_op"].CopyFrom( fetch_collection) # Run TRT optimizer in Grappler to convert the graph. self._run_conversion() self._converted_func = wrap_function.function_from_graph_def( self._converted_graph_def, [tensor.name for tensor in frozen_func.inputs], [tensor.name for tensor in frozen_func.outputs])
def h5_to_pb(h5_save_path): model = tf.keras.models.load_model(h5_save_path, compile=False) model.summary() full_model = tf.function(lambda Input: model(Input)) full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype)) # Get frozen ConcreteFunction frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() layers = [op.name for op in frozen_func.graph.get_operations()] print("-" * 50) print("Frozen model layers: ") for layer in layers: print(layer) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Save frozen graph from frozen ConcreteFunction to hard drive tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir="./frozen_models2", name="model.pb", as_text=False)
def load_graph(self, model_path): import tensorflow as tf tf.compat.v1.reset_default_graph() path_stem = os.path.dirname(model_path) if path_stem.endswith('saved_model'): imported = tf.saved_model.load(path_stem) from tensorflow.python.framework.convert_to_constants\ import convert_variables_to_constants_v2 all_sigs = imported.signatures.keys() signatures = [s for s in all_sigs if not s.startswith("_")] func = imported.signatures[signatures[0]] frozen_func = convert_variables_to_constants_v2( func, lower_control_flow=False) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) with tf.compat.v1.Session() as sess: tf.import_graph_def(graph_def, name='') return sess.graph, self.count_ops(sess.graph) else: with tf.compat.v1.Session() as sess: graph_def = tf.compat.v1.GraphDef() # print ("graph_def version:", graph_def.version) with tf.io.gfile.GFile(model_path, 'rb') as model_f: graph_def.ParseFromString(model_f.read()) tf.import_graph_def(graph_def, name='') return sess.graph, self.count_ops(sess.graph)
def from_function(func, input_names, output_names, large_model=False): if large_model: return convert_variables_to_constants_large_model(func) if get_tf_version() < LooseVersion("2.2"): frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False) else: frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False, aggressive_inlining=True) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) # output_names = [i.name for i in frozen_func.outputs] tf_reset_default_graph() with tf_session() as sess: tf.import_graph_def(graph_def, name='') input_names = inputs_without_resource(sess, input_names) graph_def = tf_optimize(input_names, output_names, graph_def) return graph_def
def testDynamicRnn(self): input_data = { "x": constant_op.constant( np.array(np.random.random_sample((3, 10, 10)), dtype=np.float32)) } cell = rnn_cell_impl.LSTMCell(10) @def_function.function(input_signature=[ tensor_spec.TensorSpec(shape=[3, 10, 10], dtype=dtypes.float32) ]) def model(x): return rnn.dynamic_rnn(cell, x, dtype=dtypes.float32) root = tracking.AutoTrackable() root.f = model input_func = root.f.get_concrete_function() output_func = convert_to_constants.convert_variables_to_constants_v2( input_func, lower_control_flow=False) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) self._testConvertedFunction(root, root.f, output_func, input_data)
def testConstructConcreteFunction(self): input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.v1 = variables.Variable(3.) root.v2 = variables.Variable(2.) root.f = def_function.function(lambda x: root.v1 * root.v2 * x) func = root.f.get_concrete_function(input_data) input_func = convert_to_constants._construct_concrete_function( func, func.graph.as_graph_def(), {}) # Test if model has enough metadata to be frozen afterwards. variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(2, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = root.f(input_data) actual_value = nest.flatten(output_func(input_data)) self.assertEqual(expected_value.numpy(), actual_value)
def get_tf_model_proto(tf_model): # define the directory for .pb model pb_model_path = "models" # define the name of .pb model pb_model_name = "mobilenet.pb" # create directory for further converted model os.makedirs(pb_model_path, exist_ok=True) # get model TF graph tf_model_graph = tf.function(lambda x: tf_model(x)) # get concrete function tf_model_graph = tf_model_graph.get_concrete_function( tf.TensorSpec(tf_model.inputs[0].shape, tf_model.inputs[0].dtype)) # obtain frozen concrete function frozen_tf_func = convert_variables_to_constants_v2(tf_model_graph) # get frozen graph frozen_tf_func.graph.as_graph_def() # save full tf model tf.io.write_graph(graph_or_graph_def=frozen_tf_func.graph, logdir=pb_model_path, name=pb_model_name, as_text=False) return os.path.join(pb_model_path, pb_model_name)
def benchmark_tensorrt_inference(self, model): # Loading TensorRT model print(f"Loading model {model}") saved_model_loaded = tf.saved_model.load( f"tensorrt_models/{model}", tags=[tag_constants.SERVING]) graph_func = saved_model_loaded.signatures[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] frozen_func = convert_variables_to_constants_v2(graph_func) inference_timings = [] print("Warmup") for i in range(50): chip, chip_id = self.chip_files[0] chip = np.expand_dims(np.array(chip), axis=0) frozen_func(tf.constant(chip.astype(np.float32))) # Run benchmark for chip, chip_id in tqdm(self.chip_files): chip = np.expand_dims(np.array(chip), axis=0) start = timer() predicted_chip = frozen_func(tf.constant(chip.astype(np.float32)))[0].numpy() end = timer() inference_timings.append(end - start) self.predictions.append((np.squeeze(predicted_chip), chip_id)) print(f"Mean: {np.mean(inference_timings)}, STD: {np.std(inference_timings)}, MEDIAN: {np.median(inference_timings)}") benchmark_summary = {"timings": inference_timings, "mean": np.mean(inference_timings), "std": np.std(inference_timings), "median": np.median(inference_timings), "90_perc": np.percentile(inference_timings, 90)} with open(f"model_scores/{model}-inference_benchmark.json", "w") as inference_json: json.dump(benchmark_summary, inference_json)
def save_and_reload(_model): with tempfile.TemporaryDirectory() as model_path: tf.saved_model.save(_model, model_path) loaded = tf.saved_model.load(model_path) func = loaded.signatures["serving_default"] frozen_func = convert_variables_to_constants_v2(func) return frozen_func
def load_graph_func(saved_model_dir: str, saved_model_tags: str, saved_model_signature_key: str): """Loads a graph function in TF2.""" imported = saved_model_load.load(export_dir=saved_model_dir, tags=saved_model_tags) graph_func = imported.signatures[saved_model_signature_key] return convert_to_constants.convert_variables_to_constants_v2(graph_func)
def model_to_graph(model: Callable) -> tf.Graph: """Convert a keras model into a wrapped function graph""" input_specs = [tf.TensorSpec(t.shape, t.dtype) for t in model.inputs] wrapped_function = tf.function(lambda x: model(x)) wrapped_function = wrapped_function.get_concrete_function(input_specs) frozen_function = cc.convert_variables_to_constants_v2(wrapped_function) return frozen_function.graph
def testKerasModel(self): input_data = constant_op.constant(1., shape=[1, 1]) # Create a simple Keras model. x = [-1, 0, 1, 2, 3, 4] y = [-3, -1, 1, 3, 5, 7] model = keras.models.Sequential( [keras.layers.Dense(units=1, input_shape=[1])]) model.compile(optimizer="sgd", loss="mean_squared_error") model.fit(x, y, epochs=1) # Get the concrete function from the Keras model. @def_function.function def to_save(x): return model(x) input_func = to_save.get_concrete_function(input_data) variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(2, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = to_save(input_data) actual_value = nest.flatten(output_func(input_data)) self.assertEqual(expected_value.numpy(), actual_value)
def export_model(ca: CAProtoModel, base_filename: str) -> None: """ Saves the weights of the model to a file beginning with base filename (this can be a path). Included will be a JSON file `base_filename.json` with the model format, topology, and weights manifest. :param ca: The model to export :param base_filename: Base file name/path at which to export it :return: None """ ca.save_weights(base_filename) # TODO: do any of these need adjusting? cf = ca.call.get_concrete_function(x=tf.TensorSpec( [None, None, None, ca.channel_n]), fire_rate=tf.constant(0.5), angle=tf.constant(0.0), step_size=tf.constant(1.0)) cf = convert_to_constants.convert_variables_to_constants_v2(cf) graph_def = cf.graph.as_graph_def() graph_json = MessageToDict(graph_def) graph_json['versions'] = dict(producer='1.14', minConsumer='1.14') model_json = { 'format': 'graph-model', 'modelTopology': graph_json, 'weightsManifest': [], } with open(base_filename + '.json', 'w') as f: json.dump(model_json, f)
def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow try: import tensorflow as tf from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 LOGGER.info( f'\n{prefix} starting export with tensorflow {tf.__version__}...') f = file.with_suffix('.pb') m = tf.function(lambda x: keras_model(x)) # full model m = m.get_concrete_function( tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) frozen_func = convert_variables_to_constants_v2(m) frozen_func.graph.as_graph_def() tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) LOGGER.info( f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') except Exception as e: LOGGER.info(f'\n{prefix} export failure: {e}')
def save_frozen_graph_tf2(self): from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # Convert Keras model to ConcreteFunction full_model = tf.function(self._model).get_concrete_function( tf.TensorSpec(self._model.inputs[0].shape, self._model.inputs[0].dtype)) # Get frozen ConcreteFunction frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() layers = [op.name for op in frozen_func.graph.get_operations()] print("-" * 50) print("Frozen model layers: ") for layer in layers: print(layer) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Save frozen graph from frozen ConcreteFunction to hard drive saved_model_folder = self.build_saved_model_folder() frozen_graph_file_name = self.build_frozen_graph_file_name() tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=saved_model_folder, name=frozen_graph_file_name, as_text=False)
def testVariableSavedModel(self): """Test a basic model with Variables with saving/loading the SavedModel.""" input_data = {"x": constant_op.constant(1., shape=[1])} root = tracking.AutoTrackable() root.v1 = variables.Variable(3.) root.v2 = variables.Variable(2.) root.f = def_function.function(lambda x: root.v1 * root.v2 * x) to_save = root.f.get_concrete_function(input_data["x"]) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save(root, save_dir, to_save) saved_model = load(save_dir) input_func = saved_model.signatures["serving_default"] variable_graph_def = input_func.graph.as_graph_def() self.assertTrue(self._hasStatefulPartitionedCallOp(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) self._testConvertedFunction(root, root.f, output_func, input_data)
def testConstSavedModel(self): """Test a basic model with functions to make sure functions are inlined.""" input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.f = def_function.function(lambda x: 2. * x) to_save = root.f.get_concrete_function(input_data) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save(root, save_dir, to_save) saved_model = load(save_dir) input_func = saved_model.signatures["serving_default"] variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(variable_graph_def)) self.assertTrue(variable_graph_def.library.function) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(constant_graph_def.library.function) # Check value. expected_value = root.f(input_data) actual_value = nest.flatten(output_func(input_data)) self.assertEqual(expected_value.numpy(), actual_value)
def _convert_saved_model_v2(self): """Convert the input SavedModel in 2.0 format.""" self._saved_model = load.load(self._input_saved_model_dir, self._input_saved_model_tags) func = self._saved_model.signatures[self._input_saved_model_signature_key] frozen_func = convert_to_constants.convert_variables_to_constants_v2(func) self._grappler_meta_graph_def = saver.export_meta_graph( graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) # Add a collection 'train_op' so that Grappler knows the outputs. fetch_collection = meta_graph_pb2.CollectionDef() for array in func.inputs + func.outputs: fetch_collection.node_list.value.append(array.name) self._grappler_meta_graph_def.collection_def["train_op"].CopyFrom( fetch_collection) # Run TRT optimizer in Grappler to convert the graph. self._run_conversion() def _get_tensor(graph, tensors): new_tensors = [] for tensor in tensors: new_tensor = graph.get_tensor_by_name(tensor.name) new_tensor.set_shape(tensor.shape) new_tensors.append(new_tensor) return new_tensors # TODO(laigd): do we need to use different name e.g. "trt_func_graph"? converted_graph = func_graph.FuncGraph(func.graph.name) with converted_graph.as_default(): importer.import_graph_def(self._converted_graph_def, name="") converted_graph.inputs = _get_tensor(converted_graph, func.graph.inputs) converted_graph.outputs = _get_tensor(converted_graph, func.graph.outputs) converted_graph.structured_outputs = func.graph.structured_outputs converted_graph.structured_input_signature = ( func.graph.structured_input_signature) # pylint: disable=protected-access # TODO(laigd): should we set up the signature as well? self._converted_func = function.ConcreteFunction( converted_graph, attrs=None, signature=None) self._converted_func.add_to_graph() self._converted_func._arg_keywords = func._arg_keywords self._converted_func._num_positional_args = func._num_positional_args self._converted_func._captured_inputs = func._captured_inputs self._converted_func.graph.variables = func.graph.variables
def testMultiFunctionModel(self): """Test a basic model with Variables.""" class BasicModel(tracking.AutoTrackable): def __init__(self): self.y = None self.z = None @def_function.function def add(self, x): if self.y is None: self.y = variables.Variable(2.) return x + self.y @def_function.function def sub(self, x): if self.z is None: self.z = variables.Variable(3.) return x - self.z input_data = constant_op.constant(1., shape=[1]) root = BasicModel() input_func = root.add.get_concrete_function(input_data) variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(1, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = root.add(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def testVariableModel(self): """Test a basic model with Variables.""" input_data = constant_op.constant(1., shape=[1]) root = tracking.AutoTrackable() root.v1 = variables.Variable(3.) root.v2 = variables.Variable(2.) root.f = def_function.function(lambda x: root.v1 * root.v2 * x) input_func = root.f.get_concrete_function(input_data) variable_graph_def = input_func.graph.as_graph_def() self.assertEqual(2, self._getNumVariables(variable_graph_def)) output_func = convert_to_constants.convert_variables_to_constants_v2( input_func) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check value. expected_value = root.f(input_data) actual_value = self._evaluateGraphDef(constant_graph_def, input_func, [input_data.numpy()]) self.assertEqual(expected_value.numpy(), actual_value)
def convert(self): """Convert the input SavedModel in 2.0 format. Returns: The TF-TRT converted Function. """ assert not self._converted self._saved_model = load.load(self._input_saved_model_dir, self._input_saved_model_tags) func = self._saved_model.signatures[self._input_saved_model_signature_key] frozen_func = convert_to_constants.convert_variables_to_constants_v2(func) grappler_meta_graph_def = saver.export_meta_graph( graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) # Add a collection 'train_op' so that Grappler knows the outputs. fetch_collection = meta_graph_pb2.CollectionDef() for array in frozen_func.inputs + frozen_func.outputs: fetch_collection.node_list.value.append(array.name) grappler_meta_graph_def.collection_def["train_op"].CopyFrom( fetch_collection) # Run TRT optimizer in Grappler to convert the graph. self._converted_graph_def = self._run_conversion(grappler_meta_graph_def) self._converted_func = wrap_function.function_from_graph_def( self._converted_graph_def, [tensor.name for tensor in frozen_func.inputs], [tensor.name for tensor in frozen_func.outputs]) self._converted = True # Wrap the converted ConcreteFunction in a Function so it can accept numpy # arrays as input. @def_function.function def wrapper_func(*args, **kwargs): return self._converted_func(*args, **kwargs) return wrapper_func
def _convert_saved_model_v2(self): """Convert the input SavedModel in 2.0 format.""" assert context.executing_eagerly() self._saved_model = load.load(self._input_saved_model_dir, self._input_saved_model_tags) func = self._saved_model.signatures[self._input_saved_model_signature_key] frozen_func = convert_to_constants.convert_variables_to_constants_v2(func) self._grappler_meta_graph_def = saver.export_meta_graph( graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) # Add a collection 'train_op' so that Grappler knows the outputs. fetch_collection = meta_graph_pb2.CollectionDef() for array in frozen_func.inputs + frozen_func.outputs: fetch_collection.node_list.value.append(array.name) self._grappler_meta_graph_def.collection_def["train_op"].CopyFrom( fetch_collection) # Run TRT optimizer in Grappler to convert the graph. self._run_conversion() self._converted_func = wrap_function.function_from_graph_def( self._converted_graph_def, [tensor.name for tensor in frozen_func.inputs], [tensor.name for tensor in frozen_func.outputs])
def convert(self): """Converts a TensorFlow GraphDef based on instance variables. Returns: The converted data in serialized format. Raises: ValueError: Input shape is not specified. None value for dimension in input_tensor. """ frozen_func = _convert_to_constants.convert_variables_to_constants_v2( self._func) input_tensors = [ tensor for tensor in frozen_func.inputs if tensor.dtype != _dtypes.resource ] output_tensors = frozen_func.outputs # Run a Grappler pass. is_only_flex_enabled = set( [OpsSet.SELECT_TF_OPS]) == self.target_spec.supported_ops config = _get_grappler_config(enable_layout_optimizer=is_only_flex_enabled) graph_def = _run_graph_optimizations( frozen_func.graph.as_graph_def(), input_tensors, output_tensors, 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) if self.representative_dataset: if not isinstance(self.representative_dataset, RepresentativeDataset): raise TypeError("`representative_dataset` must be an instance of " "`RepresentativeDataset`") if self.representative_dataset.input_gen is None: raise ValueError( "Provide an input generator for `representative_dataset`") # TODO(shashishekhar): For now use optimizations order is ignored. # Both size and latency optimizations decide whether to apply post # training optimizations. post_training_optimize = bool( len( set(self.optimizations) & set([Optimize.OPTIMIZE_FOR_LATENCY, Optimize.OPTIMIZE_FOR_SIZE]))) # Do weights only quantization if there is no dataset for calibration. weights_only_quantize_flag = ( post_training_optimize and (self.representative_dataset is None)) converter_kwargs = { "input_format": constants.TENSORFLOW_GRAPHDEF, "allow_custom_ops": self.allow_custom_ops, "post_training_quantize": weights_only_quantize_flag, "target_ops": self.target_spec.supported_ops, } # Converts model. result = _toco_convert_impl( input_data=graph_def, input_tensors=input_tensors, output_tensors=output_tensors, **converter_kwargs) if self.representative_dataset and post_training_optimize: calibrate_quantize = _calibrator.Calibrator(result) result = calibrate_quantize.calibrate_and_quantize( self.representative_dataset.input_gen) return result
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. self._target_ops = self.target_spec.supported_ops if len(self._funcs) != 1: raise ValueError("This converter can only convert a single " "ConcreteFunction. Converting multiple functions is " "under development.") frozen_func = _convert_to_constants.convert_variables_to_constants_v2( self._funcs[0]) 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 = frozen_func.graph.as_graph_def() 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_representative_dataset() converter_kwargs = { "input_format": constants.TENSORFLOW_GRAPHDEF, "allow_custom_ops": self.allow_custom_ops, "post_training_quantize": self._is_weight_only_quantize(), "target_ops": self.target_spec.supported_ops, } # 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) return result
def from_keras_model_file(cls, model_file, input_arrays=None, input_shapes=None, output_arrays=None, custom_objects=None): """Creates a TFLiteConverter class from a tf.keras model file. Args: model_file: Full filepath of HDF5 file containing the tf.keras model. input_arrays: List of input tensors to freeze graph with. Uses input arrays from SignatureDef when none are provided. (default None) input_shapes: Dict of strings representing input tensor names to list of integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}). Automatically determined when input shapes is None (e.g., {"foo" : None}). (default None) output_arrays: List of output tensors to freeze graph with. Uses output arrays from SignatureDef when none are provided. (default None) custom_objects: Dict mapping names (strings) to custom classes or functions to be considered during model deserialization. (default None) Returns: TFLiteConverter class. """ # Handles Keras when Eager mode is enabled. if context.executing_eagerly(): if input_arrays or output_arrays: raise ValueError("`input_arrays` and `output_arrays` are unsupported " "with Eager mode. If your model requires any of these " "parameters, please use disable_eager_execution().") _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) function = _saving_utils.trace_model_call(keras_model) concrete_func = function.get_concrete_function() frozen_func = _convert_to_constants.convert_variables_to_constants_v2( concrete_func) _set_tensor_shapes(frozen_func.inputs, input_shapes) return cls(frozen_func.graph.as_graph_def(), frozen_func.inputs, frozen_func.outputs) # Handles Keras when Eager mode is disabled. _keras.backend.clear_session() _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) sess = _keras.backend.get_session() # Get input and output tensors. if input_arrays: input_tensors = _get_tensors_from_tensor_names(sess.graph, input_arrays) else: input_tensors = keras_model.inputs if output_arrays: output_tensors = _get_tensors_from_tensor_names(sess.graph, output_arrays) else: output_tensors = keras_model.outputs _set_tensor_shapes(input_tensors, input_shapes) graph_def = _freeze_graph(sess, input_tensors, output_tensors) return cls(graph_def, input_tensors, output_tensors)