def test_asset_loading(self): first_path = self._v1_asset_saved_model() imported = load.load(first_path) self.evaluate(lookup_ops.tables_initializer()) fn = imported.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) second_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(imported, second_path, signatures=imported.signatures) shutil.rmtree(first_path) del ops.get_collection_ref(ops.GraphKeys.TABLE_INITIALIZERS)[:] second_import = load.load(second_path) self.evaluate(lookup_ops.tables_initializer()) fn = second_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) third_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(second_import, third_path, signatures=second_import.signatures) shutil.rmtree(second_path) del ops.get_collection_ref(ops.GraphKeys.TABLE_INITIALIZERS)[:] third_import = load.load(third_path) self.evaluate(lookup_ops.tables_initializer()) fn = third_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"])))
def test_trainable_not_set_in_proto(self): """If a VariableDef has no 'trainable', we fall back to collections.""" real_tf_version = versions.__version__ # Pretend to be exported from an older version of TensorFlow, so trainable # will follow collections instead of checking VariableDefs. versions.__version__ = "1.7.0" path = self._no_trainable_variable_attribute(trainable=True) root = load.load(path) self.assertTrue(root.variables[0].trainable) path = self._no_trainable_variable_attribute(trainable=False) root = load.load(path) self.assertFalse(root.variables[0].trainable) versions.__version__ = real_tf_version
def test_trainable_in_proto(self): """If a VariableDef has a trainable property, we do not use collections.""" path = self._export_variable( trainable=True, collections=[ops.GraphKeys.GLOBAL_VARIABLES]) root = load.load(path) self.assertTrue(root.variables[0].trainable) path = self._export_variable( trainable=False, collections=[ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.TRAINABLE_VARIABLES]) root = load.load(path) self.assertFalse(root.variables[0].trainable)
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 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 test_saved_model_v2(directory, tag_set=None, signature_key=None, input_data=None, **kwargs): """Validates the TensorFlow SavedModel converts to a TFLite model. Converts the TensorFlow SavedModel to TFLite and checks the accuracy of the model on random data. Args: directory: SavedModel directory to convert. tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to analyze. All tags in the tag set must be present. signature_key: Key identifying SignatureDef containing inputs and outputs. input_data: np.ndarray to pass into models during inference. (default None) **kwargs: Additional arguments to be passed into the converter. """ model = _load.load(directory, tags=tag_set) if not signature_key: signature_key = _signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY concrete_func = model.signatures[signature_key] converter = _lite.TFLiteConverterV2.from_concrete_functions([concrete_func]) tflite_model = _convert(converter, version=2, **kwargs) compare_models_v2(tflite_model, concrete_func, input_data=input_data)
def _TestTrtGraphConverter(self, input_saved_model_dir=None, output_saved_model_dir=None, need_calibration=False, is_dynamic_op=False): """General method to test trt_convert.TrtGraphConverter().""" output_graph_def = self._ConvertGraph( input_saved_model_dir=input_saved_model_dir, output_saved_model_dir=output_saved_model_dir, need_calibration=need_calibration, is_dynamic_op=is_dynamic_op, use_function_backup=need_calibration) graph_defs_to_verify = [output_graph_def] if output_saved_model_dir: if context.executing_eagerly(): root = load.load(output_saved_model_dir) saved_model_graph_def = root.signatures[ signature_constants .DEFAULT_SERVING_SIGNATURE_DEF_KEY].graph.as_graph_def() else: saved_model_graph_def = saved_model_utils.get_meta_graph_def( output_saved_model_dir, tag_constants.SERVING).graph_def self.assertTrue(isinstance(saved_model_graph_def, graph_pb2.GraphDef)) graph_defs_to_verify.append(saved_model_graph_def) for graph_def in graph_defs_to_verify: node_name_to_op = {node.name: node.op for node in graph_def.node} if context.executing_eagerly(): # In V2 the actual graph could be inside a function. for func in graph_def.library.function: node_name_to_op.update({node.name: node.op for node in func.node_def}) self.assertIn("TRTEngineOp_0", node_name_to_op) self.assertEqual("TRTEngineOp", node_name_to_op["TRTEngineOp_0"]) else: self.assertEqual({ "input": "Placeholder", "TRTEngineOp_0": "TRTEngineOp", "output": "Identity" }, node_name_to_op) if need_calibration: trt_engine_nodes = [ node for node in graph_def.node if node.op == "TRTEngineOp" ] self.assertNotEmpty(trt_engine_nodes) for node in trt_engine_nodes: self.assertTrue(len(node.attr["calibration_data"].s)) # Run the calibrated graph. # TODO(laigd): consider having some input where the answer is different. with ops.Graph().as_default(): importer.import_graph_def(graph_def, name="") with self.session(config=self._GetConfigProto()) as sess: for test_data in range(10): self.assertEqual((test_data + 1.0)**2, sess.run( "output:0", feed_dict={"input:0": [[[test_data]]]}))
def test_multi_meta_graph_loading(self): with self.assertRaisesRegexp(ValueError, "2 MetaGraphs"): load.load(self._v1_multi_metagraph_saved_model()) first_imported = load.load(self._v1_multi_metagraph_saved_model(), tags=["first"]) self.assertEqual({"first_output": 42.}, self.evaluate(first_imported.signatures["first_key"]( first_start=constant_op.constant(2.)))) second_imported = load.load(self._v1_multi_metagraph_saved_model(), tags=set(["second"])) with self.assertRaisesRegexp(TypeError, "second_start"): second_imported.signatures["second_key"](x=constant_op.constant(2.)) with self.assertRaisesRegexp(TypeError, "second_start"): second_imported.signatures["second_key"]( second_start=constant_op.constant(2.), x=constant_op.constant(2.)) self.assertEqual({"second_output": 21.}, self.evaluate(second_imported.signatures["second_key"]( second_start=constant_op.constant(2.))))
def cycle(self, obj, cycles=1, signatures=None): to_save = obj # TODO(vbardiovsky): It would be nice if exported protos reached a fixed # point w.r.t. saving/restoring, ideally after 2nd saving. for _ in range(cycles): path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(to_save, path, signatures) loaded = load.load(path) to_save = loaded return loaded
def cycle(obj, cycles, signatures=None, options=None): to_save = obj for _ in range(cycles): path = tempfile.mkdtemp(prefix=test.get_temp_dir()) # If available, we'll run the save and restore preferring the GPU. This # just makes sure we aren't throwing errors and have enough # device("CPU") blocks to satisfy the placer. with test_util.use_gpu(): save.save(to_save, path, signatures, options=options) loaded = load.load(path) signatures = loaded.signatures to_save = loaded return loaded
def testSaveAndLoadModuleUnderStrategy(self): class Dense(module.Module): def __init__(self): self.kernel = variables_lib.Variable(random_ops.random_uniform( (6, 6)), name='kernel') self.bias = variables_lib.Variable(random_ops.random_uniform( (6, )), name='bias') @def_function.function def __call__(self, x): out = math_ops.matmul(self.kernel, x) out = out + self.bias return out x = constant_op.constant(math_ops.range(6, dtype=dtypes.float32), shape=[6, 1]) strategy = self._create_strategy(2) with strategy.scope(): layer = Dense() expect = layer(x) model_dir = self.get_temp_dir() save.save(layer, model_dir) strategy2 = self._create_strategy(3) with strategy2.scope(): loaded_layer = load.load(model_dir) # Should fail with informative error with self.assertRaisesRegex(ValueError, 'run a loaded non-Keras'): got = loaded_layer(x) # Loading without a strategy should work, because the tf.function is traced # with a single variable as input loaded_layer = load.load(model_dir) got = loaded_layer(x) self.assertAllClose(got, expect)
def test_load_with_tags(self): root = tracking.AutoTrackable() path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(root, path) with self.assertRaises(ValueError): load.load(path, tags=[tag_constants.EVAL]) load.load(path, tags=[tag_constants.SERVING]) load.load(path, tags=tag_constants.SERVING)
def testLoadSavedModelWithUnregisteredExtensionType(self): def f(x, y): x_values = x.values if isinstance(x, MaskedTensorV1) else x y_values = y.values if isinstance(y, MaskedTensorV1) else y x_mask = x.mask if isinstance(x, MaskedTensorV1) else True y_mask = y.mask if isinstance(y, MaskedTensorV1) else True return MaskedTensorV1(x_values + y_values, x_mask & y_mask) t_spec = tensor_spec.TensorSpec(None, dtypes.int32) b_spec = tensor_spec.TensorSpec(None, dtypes.bool) mt_spec = MaskedTensorV1.Spec(values=t_spec, mask=b_spec) model = module.Module() model.f = def_function.function(f) model.f.get_concrete_function(t_spec, t_spec) model.f.get_concrete_function(t_spec, mt_spec) model.f.get_concrete_function(mt_spec, t_spec) model.f.get_concrete_function(mt_spec, mt_spec) path = tempfile.mkdtemp(prefix=test.get_temp_dir()) with temporarily_register_type_spec('tf.test.MaskedTensorV1.Spec', MaskedTensorV1.Spec): save.save(model, path) loaded_model = load.load(path) with self.assertRaises(ValueError): type_spec.lookup('tf.test.MaskedTensorV1') t = constant_op.constant([10, 20, 30]) v1 = loaded_model.f(t, t) self.assertIsInstance(v1, extension_type.AnonymousExtensionType) self.assertAllEqual(v1.values, [20, 40, 60]) self.assertAllEqual(v1.mask, True) v2 = loaded_model.f(v1, v1) self.assertIsInstance(v2, extension_type.AnonymousExtensionType) self.assertAllEqual(v2.values, [40, 80, 120]) self.assertAllEqual(v2.mask, True) mt = MaskedTensorV1([1, 2, 3], [True, True, False]) v3 = loaded_model.f( t, extension_type.reinterpret(mt, extension_type.AnonymousExtensionType)) self.assertIsInstance(v3, extension_type.AnonymousExtensionType) self.assertAllEqual(v3.values, [11, 22, 33]) self.assertAllEqual(v3.mask, [True, True, False]) v4 = extension_type.reinterpret(v3, MaskedTensorV1) self.assertIsInstance(v4, MaskedTensorV1) self.assertAllEqual(v4.values, [11, 22, 33]) self.assertAllEqual(v4.mask, [True, True, False])
def test_saved_model(self): with self.device: different_values = self.device.pack( [constant_op.constant(-1.), constant_op.constant(3.)]) m = module.Module() m.v = variables.Variable(different_values) m.f = def_function.function(lambda: m.v * 2.) self.assertAllClose([-2., 6.], self.device.unpack(m.f())) saved_model_path = os.path.join(self.get_temp_dir(), "saved_model") save.save(m, saved_model_path) context._reset_context() self.setUp() single_device_loaded = load.load(saved_model_path) self.assertAllClose(-2., single_device_loaded.f()) with self.device: parallel_loaded = load.load(saved_model_path) self.assertAllClose([-2., 6.], self.device.unpack(parallel_loaded.f())) self.assertAllClose([-1., 3.], self.device.unpack(parallel_loaded.v)) parallel_loaded.v.assign(self.device.pack([.1, .2])) self.assertAllClose([.2, .4], self.device.unpack(parallel_loaded.f()))
def testRefVariableImport(self): saved = self._singleMetaGraphSavedModel() imported = load(saved) fn = imported.signatures["serving_default"] output_func = convert_to_constants.convert_variables_to_constants_v2( fn) constant_graph_def = output_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) input_data = {"start": constant_op.constant(1., shape=[1, 1])} root = tracking.AutoTrackable() self._testConvertedFunction(root, fn, output_func, input_data)
def convert_tf_saved_model(saved_model_dir, output_dir, signature_def='serving_default', saved_model_tags='serve', quantization_dtype=None, skip_op_check=False, strip_debug_ops=False): """Freeze the SavedModel and check the model compatibility with Tensorflow.js. Optimize and convert the model to Tensorflow.js format, when the model passes the compatiblity check. Args: saved_model_dir: string The saved model directory. : string The names of the output nodes, comma separated. output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. signature_def: string Tagset of the SignatureDef to load. Defaults to 'serving_default'. saved_model_tags: tags of the GraphDef to load. Defaults to 'serve'. quantization_dtype: An optional numpy dtype to quantize weights to for compression. Only np.uint8 and np.uint16 are supported. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. """ if signature_def is None: signature_def = 'serving_default' if not os.path.exists(output_dir): os.makedirs(output_dir) output_graph = os.path.join(output_dir, common.ARTIFACT_MODEL_JSON_FILE_NAME) saved_model_tags = saved_model_tags.split(', ') model = load(saved_model_dir, saved_model_tags) _check_signature_in_model(model, signature_def) concrete_func = model.signatures[signature_def] frozen_func = convert_to_constants.convert_variables_to_constants_v2( concrete_func) optimize_graph(frozen_func, output_graph, model.tensorflow_version, quantization_dtype=quantization_dtype, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops)
def test_structure_import(self): root = tracking.Checkpointable() root.f = def_function.function( lambda x: 2. * x, input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) root.dep_one = tracking.Checkpointable() root.dep_two = tracking.Checkpointable() root.dep_two.dep = tracking.Checkpointable() root.dep_three = root.dep_two.dep save_dir = os.path.join(self.get_temp_dir(), "saved_model") save.save(root, save_dir) imported = load.load(save_dir) self.assertIs(imported.dep_three, imported.dep_two.dep) self.assertIsNot(imported.dep_one, imported.dep_two)
def test_asset_loading(self): first_path = self._v1_asset_saved_model() imported = load.load(first_path) fn = imported.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) second_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(imported, second_path, signatures=imported.signatures) shutil.rmtree(first_path) second_import = load.load(second_path) fn = second_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) third_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(second_import, third_path, signatures=second_import.signatures) shutil.rmtree(second_path) third_import = load.load(third_path) fn = third_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"])))
def test_capture_assets(self): root = tracking.Checkpointable() root.vocab = tracking.TrackableAsset(self._make_asset("contents")) root.f = def_function.function( lambda: root.vocab.asset_path, input_signature=[]) save_dir = os.path.join(self.get_temp_dir(), "save_dir") save.save(root, save_dir) imported = load.load(save_dir) origin_output = root.f().numpy() imported_output = imported.f().numpy() self.assertNotEqual(origin_output, imported_output) with open(imported_output, "r") as f: self.assertEquals("contents", f.read())
def test_asset_loading(self): first_path = self._v1_asset_saved_model() imported = load.load(first_path) fn = imported.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) second_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(imported, second_path, signatures=imported.signatures) shutil.rmtree(first_path) second_import = load.load(second_path) fn = second_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))) third_path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid())) save.save(second_import, third_path, signatures=second_import.signatures) shutil.rmtree(second_path) third_import = load.load(third_path) fn = third_import.signatures["serving_default"] self.assertAllClose({"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"])))
def testTrtGraphConverter_Int8Conversion_v2(self): if not is_tensorrt_enabled(): return np_input = np.random.random_sample([4, 1, 1]).astype(np.float32) # Create a model and save it. input_saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir()) root = self._GetModelForV2() expected_output = root.run(np_input) save.save(root, input_saved_model_dir, {_SAVED_MODEL_SIGNATURE_KEY: root.run}) # Run TRT conversion. converter = self._CreateConverterV2( input_saved_model_dir, precision_mode=trt_convert.TrtPrecisionMode.INT8) converted_func = converter.convert() # Run the converted function for INT8 calibration. calibration_output = converted_func(np_input) self.assertEqual(1, len(calibration_output)) self.assertAllClose(expected_output, calibration_output.values()[0], atol=1e-6, rtol=1e-6) # Save the converted model again with serialized engine cache. output_saved_model_dir = self.mkdtemp() converter.save(output_saved_model_dir) expected_asset_file = os.path.join( output_saved_model_dir, "assets/trt-serialized-engine.TRTEngineOp_0") self.assertTrue(os.path.exists(expected_asset_file)) self.assertTrue(os.path.getsize(expected_asset_file)) # Load and verify the converted model. root_with_trt = load.load(output_saved_model_dir) converted_signature = root_with_trt.signatures[ _SAVED_MODEL_SIGNATURE_KEY] output_with_trt = converted_signature(ops.convert_to_tensor(np_input)) self.assertEqual(1, len(output_with_trt)) # The output of running the converted signature is a dict due to # compatibility reasons with V1 SavedModel signature mechanism. self.assertAllClose(expected_output, output_with_trt.values()[0], atol=1e-6, rtol=1e-6)
def testTrtGraphConverter_ShapeOp_v2(self): """Test case for TrtGraphConverterV2 with ShapeOp.""" class ShapeOpModel(tracking.AutoTrackable): def __init__(self): self.v = None @def_function.function(input_signature=[ tensor_spec.TensorSpec(shape=[None, None], dtype=dtypes.float32) ]) def run(self, x): q = x + 1 q_shape = array_ops.shape(q) return array_ops.identity(q_shape, name="output") np_input = np.random.random_sample([5, 3]).astype(np.float32) def _InputFunc(): yield (np_input,) # Create the SavedModel. root = ShapeOpModel() expected_output = root.run(np_input) input_saved_model_dir = self.mkdtemp() save.save(root, input_saved_model_dir, signatures=root.run) # Convert the graph to TF-TRT. conv_params = trt_convert.TrtConversionParams(minimum_segment_size=2) converter = trt_convert.TrtGraphConverterV2( input_saved_model_dir=input_saved_model_dir, conversion_params=conv_params, use_dynamic_shape=True) converter.convert() # Build the graph with the input generator. This runs the TRTEngineOp native # segment. converter.build(_InputFunc) output_saved_model_dir = self.mkdtemp() converter.save(output_saved_model_dir) root_with_trt = load.load(output_saved_model_dir) converted_signature = root_with_trt.signatures["serving_default"] # Check that the graph is converted to one TRTEngineOp. self._CheckTrtOps(converted_signature) # Run the graph. output_with_trt = converted_signature(x=ops.convert_to_tensor(np_input)) # Check the result of the run. self.assertAllClose(expected_output, list(output_with_trt.values())[0])
def testTrainingDefaults(self): def assert_training_default(fn, default_value): arg_spec = tf_inspect.getfullargspec(fn) index = len(arg_spec.args) - arg_spec.args.index('training') self.assertEqual(arg_spec.defaults[-index], default_value) class LayerWithTrainingRequiredArg(keras.engine.base_layer.Layer): def call(self, inputs, training): return tf_utils.smart_cond(training, lambda: inputs * 0, lambda: array_ops.identity(inputs)) class LayerWithTrainingDefaultTrue(keras.engine.base_layer.Layer): def call(self, inputs, training=True): return tf_utils.smart_cond(training, lambda: inputs * 0, lambda: array_ops.identity(inputs)) class Model(keras.models.Model): def __init__(self): super(Model, self).__init__() self.layer_with_training_default_none = LayerWithLearningPhase( ) self.layer_with_training_default_true = LayerWithTrainingDefaultTrue( ) self.layer_with_required_training_arg = LayerWithTrainingRequiredArg( ) def call(self, inputs): x = self.layer_with_training_default_none(inputs) x += self.layer_with_training_default_true(inputs) x += self.layer_with_required_training_arg(inputs, False) return x model = Model() # Build and set model inputs model.predict(np.ones([1, 3]).astype('float32')) saved_model_dir = self._save_model_dir() model.save(saved_model_dir, save_format='tf') load = tf_load.load(saved_model_dir) assert_training_default(load.__call__, False) assert_training_default(load.layer_with_training_default_none.__call__, False) assert_training_default(load.layer_with_training_default_true.__call__, True) # Assert that there are no defaults for layer with required training arg arg_spec = tf_inspect.getfullargspec( load.layer_with_required_training_arg.__call__) self.assertFalse(arg_spec.defaults) # defaults is None or empty
def test_save_uninitialized_variable(self): root = tracking.AutoTrackable() root.uninitialized_variable = resource_variable_ops.UninitializedVariable( name="uninitialized_variable", dtype=dtypes.float32) root.initialized_variable = variables.Variable( 1.0, name="initialized_variable") # TODO(b/149594077): Python loading does not work now partly because it # shouldn't, as the public API and semantics of uninitialized variables # are not properly defined, and officially supporting loading would end up # defining semantics "by usage." We should only allow loading once the API # is made official. export_dir = os.path.join(self.get_temp_dir(), "saved_model") save.save(root, export_dir) with self.assertRaisesRegex(FileNotFoundError, "Key uninitialized_variable"): load.load(export_dir) with ops.Graph().as_default(), session_lib.Session() as session: # The final ValueError here (with "no variables to save") is confusing, # but errors upstream give the user the correct information (a # NotFoundError stating that the uninitalized_variable was not found in # the checkpoint). with self.assertRaises(ValueError): loader.load(session, [tag_constants.SERVING], export_dir)
def test_concrete_function_with_set_shape(self,): # Serialized concrete function should retain the shape from the TensorSpec, # instead of using the shape of the inputs (which are changed by set_shape). @def_function.function def f(x): x.set_shape((5, 1)) return x root = tracking.AutoTrackable() path = os.path.join(self.get_temp_dir(), "saved_model") concrete = f.get_concrete_function( tensor_spec.TensorSpec((None, 1), name="name")) save.save(root, path, signatures={"key": concrete}) imported = load.load(path) self.assertEqual(imported.signatures["key"].structured_input_signature[1], {"name": tensor_spec.TensorSpec((None, 1), name="name")})
def _testConvertedFunction(self, obj, func, converted_concrete_func, input_data): # Check that the converted ConcreteFunction produces the same result as the # original Function. expected_value = func(input_data) actual_value = nest.flatten(converted_concrete_func(input_data)) self.assertEqual(expected_value.numpy(), actual_value) # Save the converted ConcreteFunction as a signature. save_dir = os.path.join(self.get_temp_dir(), "frozen_saved_model") save(obj, save_dir, {"mykey": converted_concrete_func}) # Load it back and make sure it works. loaded_obj = load(save_dir) actual_value = nest.flatten(loaded_obj.signatures["mykey"](input_data)) self.assertEqual(expected_value.numpy(), actual_value)
def test_assets_dedup(self): vocab = self._make_asset("contents") root = tracking.Checkpointable() root.f = def_function.function( lambda x: 2. * x, input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) root.asset1 = tracking.TrackableAsset(vocab) root.asset2 = tracking.TrackableAsset(vocab) export_dir = os.path.join(self.get_temp_dir(), "save_dir") save.save(root, export_dir) imported = load.load(export_dir) self.assertEqual(imported.asset1.asset_path.numpy(), imported.asset2.asset_path.numpy())
def testBackwardCompatibility(self): """Load and execute a model that was saved in TF2.0.""" model_dir = test.test_src_dir_path( "python/compiler/tensorrt/test/testdata/tftrt_2.0_saved_model") saved_model_loaded = load.load(model_dir, tags=[tag_constants.SERVING]) graph_func = saved_model_loaded.signatures[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] np_input1 = ops.convert_to_tensor(np.ones([4, 1, 1]).astype(np.float32)) np_input2 = ops.convert_to_tensor(np.ones([4, 1, 1]).astype(np.float32)) output = graph_func(input1=np_input1, input2=np_input2)["output_0"] self.assertEqual(output.shape, (4, 1, 1)) self.assertAllClose( np.asarray([5.0, 5.0, 5.0, 5.0]).reshape([4, 1, 1]), output)
def _show_defined_functions(saved_model_dir): """Prints the callable concrete and polymorphic functions of the Saved Model. Args: saved_model_dir: Directory containing the SavedModel to inspect. """ meta_graphs = saved_model_utils.read_saved_model( saved_model_dir).meta_graphs has_object_graph_def = False for meta_graph_def in meta_graphs: has_object_graph_def |= meta_graph_def.HasField('object_graph_def') if not has_object_graph_def: return with ops_lib.Graph().as_default(): trackable_object = load.load(saved_model_dir) print('\nDefined Functions:', end='') children = list( save._AugmentedGraphView(trackable_object) # pylint: disable=protected-access .list_children(trackable_object)) children = sorted(children, key=lambda x: x.name) for name, child in children: concrete_functions = [] if isinstance(child, defun.ConcreteFunction): concrete_functions.append(child) elif isinstance(child, def_function.Function): concrete_functions.extend( child._list_all_concrete_functions_for_serialization()) # pylint: disable=protected-access else: continue print('\n Function Name: \'%s\'' % name) concrete_functions = sorted(concrete_functions, key=lambda x: x.name) for index, concrete_function in enumerate(concrete_functions, 1): args, kwargs = None, None if concrete_function.structured_input_signature: args, kwargs = concrete_function.structured_input_signature elif concrete_function._arg_keywords: # pylint: disable=protected-access # For pure ConcreteFunctions we might have nothing better than # _arg_keywords. args = concrete_function._arg_keywords # pylint: disable=protected-access if args: print(' Option #%d' % index) print(' Callable with:') _print_args(args, indent=4) if kwargs: _print_args(kwargs, 'Named Argument', indent=4)
def test_save_composite_tensor_signature(self): @def_function.function( input_signature=[ragged_tensor.RaggedTensorSpec(ragged_rank=2)]) def f(x): return {"output_key": x} root = tracking.AutoTrackable() path = os.path.join(self.get_temp_dir(), "saved_model") inp = ragged_factory_ops.constant([[[1.0, 2.0], [3.0]], [[5.]]]) flat_inp = { "x": constant_op.constant([1., 2., 3., 5]), "x_1": constant_op.constant([0, 2, 3], dtype=dtypes.int64), "x_2": constant_op.constant([0, 2, 3, 4], dtype=dtypes.int64) } save.save(root, path, signatures={"key": f.get_concrete_function()}) # Test that the ragged signature can be loaded back into Python with V2 APIs imported = load.load(path) self.assertAllEqual( inp, imported.signatures["key"](**flat_inp)["output_key"]) graph = ops.Graph() # Try running the signature with V1 APIs. with graph.as_default(), session_lib.Session() as session: meta_graph_def = loader.load(session, [tag_constants.SERVING], path) signature = meta_graph_def.signature_def["key"] feed_dict = {} for arg_name in flat_inp: input_tensor = session.graph.get_tensor_by_name( signature.inputs[arg_name].name) feed_dict[input_tensor] = flat_inp[arg_name].numpy() # Get composite tensor components output_components = ( signature.outputs["output_key"].composite_tensor.components) fetches = {} components_keys = ["x", "x_1", "x_2"] for k, output_tensor_info in zip(components_keys, output_components): fetches[k] = session.graph.get_tensor_by_name( output_tensor_info.name) outputs = session.run(fetches, feed_dict) self.assertAllClose(flat_inp, outputs)
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 _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 test_only_implicit_signatures(self): def func(x): return 2 * x root = tracking.Checkpointable() root.f = def_function.function(func) # Add two traces. root.f(constant_op.constant(1.)) root.f(constant_op.constant(1)) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save.save(root, save_dir, signatures=dict()) imported = load.load(save_dir) self.assertEqual(4., imported.f(constant_op.constant(2.)).numpy()) self.assertEqual(14, imported.f(constant_op.constant(7)).numpy())
def test_table_in_graph(self, cycles): root = self._make_model_with_tables() if cycles > 1: root = self.cycle(root, cycles - 1) path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(root, path) imported = self.cycle(root, 1) with ops.Graph().as_default(): imported = load.load(path) keys = constant_op.constant(["brain", "test", "foo", "surgery"]) output1 = imported.lookup1(keys) output2 = imported.lookup2(keys) with monitored_session.MonitoredSession() as sess: self.assertAllEqual([0, -1, -1, 2], sess.run(output1)) self.assertAllEqual([2, 0, 1, -1], sess.run(output2))
def test_save_without_tracing(self): class DoNotTrace(keras.layers.Layer): def __init__(self): super(DoNotTrace, self).__init__() self.input_spec = keras.layers.InputSpec(shape=[None]) self.built = True def call(self, inputs): raise ValueError('I said do not trace') def get_config(self): return {} @property def _use_input_spec_as_call_signature(self): return True root = keras.models.Sequential() root.add(keras.layers.Input(shape=(3,))) root.attached_layer = DoNotTrace() saved_model_dir = self._save_model_dir() # With the default settings, the call function is traced. with self.assertRaisesRegex(ValueError, 'do not trace'): root.save(saved_model_dir, save_format='tf') # When saving the config only, the layer call function should not be not # traced. root.save(saved_model_dir, save_format='tf', save_traces=False) loaded = tf_load.load(saved_model_dir) self.assertTrue(hasattr(loaded, 'attached_layer')) # This should raise an error when loaded without the custom object loaded = keras_load.load(saved_model_dir) with self.assertRaisesRegex(ValueError, 'Cannot call custom layer'): loaded.attached_layer(constant_op.constant([1.])) # Try loading with the custom objects with generic_utils.CustomObjectScope({'DoNotTrace': DoNotTrace}): loaded = keras_load.load(saved_model_dir) with self.assertRaisesRegex(ValueError, 'I said do not trace'): loaded.attached_layer(constant_op.constant([1.]))
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]) # Reconstruct the output signatures using the ones from original model. self._converted_func.graph.structured_outputs = nest.pack_sequence_as( func.graph.structured_outputs, self._converted_func.graph.structured_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 testConstSavedModel(self): """Test a basic model with functions to make sure functions are inlined.""" input_data = {"x": 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["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.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) self._testConvertedFunction(root, root.f, output_func, input_data)
def test_resource_variable_import(self): imported = load.load(self._v1_single_metagraph_saved_model( use_resource=True)) fn = imported.signatures["serving_default"] self.assertEqual({"output": 6.}, self.evaluate(fn(constant_op.constant(2.)))) self.assertAllEqual([3., 1.], self.evaluate(imported.variables)) imported.variables[0].assign(4.) self.assertEqual({"output": 8.}, self.evaluate(fn(start=constant_op.constant(2.)))) imported.variables[1].assign(2.) self.assertEqual({"output": 24.}, self.evaluate(fn(start=constant_op.constant(3.)))) self.assertTrue(imported.variables[0].trainable) self.assertFalse(imported.variables[1].trainable) with backprop.GradientTape() as tape: output = fn(start=constant_op.constant(4.)) self.assertEqual(imported.variables[:1], list(tape.watched_variables())) self.assertEqual(8., tape.gradient(output, imported.variables[0]).numpy())
def testSavedModelSupport(self): class TestModule(module.Module): @def_function.function def f(self, s): return s.x[0] + s.x[1] + s.y s1 = self.StructWithName((1, 2), 3) s2 = self.StructWithName((1.0, 2), [3.0, 4.0]) m = TestModule() m.f.get_concrete_function(s1) m.f.get_concrete_function(s2) path = tempfile.mkdtemp(prefix=test.get_temp_dir()) save.save(m, path) loaded = load.load(path) self.assertAllEqual(loaded.f(s1), 6) self.assertAllEqual(loaded.f(s2), [6.0, 7.0])
def _GetGraphDef(self, run_params, gdef_or_saved_model_dir): if isinstance(gdef_or_saved_model_dir, str): if run_params.is_v2: root = load.load(gdef_or_saved_model_dir) func = root.signatures[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] gdef = func.graph.as_graph_def() # Manually unref the loaded saved model and force GC to destroy the TRT # engine cache after load(). There is currently a reference cycle in 2.0 # which prevents auto deletion of the resource. # TODO(laigd): fix this. del func del root gc.collect() return gdef return saved_model_utils.get_meta_graph_def( gdef_or_saved_model_dir, tag_constants.SERVING).graph_def assert isinstance(gdef_or_saved_model_dir, graph_pb2.GraphDef) return gdef_or_saved_model_dir
def test_load_in_graph_mode(self, cycles): root = tracking.AutoCheckpointable() root.v1 = variables.Variable(1.) root.v2 = variables.Variable(2.) root.f = def_function.function( lambda x: root.v2 * x, input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) if cycles > 1: root = self.cycle(root, cycles - 1) path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(root, path) with ops.Graph().as_default(): imported = load.load(path) var_v1 = imported.v1 output = imported.f(constant_op.constant(2.)) with monitored_session.MonitoredSession() as sess: self.assertEqual(1.0, sess.run(var_v1)) self.assertEqual(4.0, sess.run(output))
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) self._testConvertedFunction(root, root.f, output_func, input_data)
def test_load_in_graph_mode(self, cycles): root = tracking.AutoTrackable() root.v1 = variables.Variable(1.) root.v2 = variables.Variable(2.) root.f = def_function.function( lambda x: root.v2 * x, input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) if cycles > 1: root = self.cycle(root, cycles - 1) path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(root, path) with ops.Graph().as_default(): imported = load.load(path) var_v1 = imported.v1 output = imported.f(constant_op.constant(2.)) with monitored_session.MonitoredSession() as sess: self.assertEqual(1.0, sess.run(var_v1)) self.assertEqual(4.0, sess.run(output))
def testConstSavedModel(self): """Test a basic model with functions to make sure functions are inlined.""" self.skipTest('b/124205572') 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) concrete_func = saved_model.signatures['serving_default'] # Convert model and ensure model is not None. converter = lite.TFLiteConverterV2.from_concrete_function(concrete_func) tflite_model = converter.convert() # Check values from converted model. expected_value = root.f(input_data) actual_value = self._evaluateTFLiteModel(tflite_model, [input_data]) self.assertEqual(expected_value.numpy(), actual_value)
def test_assets(self): file1 = self._make_asset("contents 1") file2 = self._make_asset("contents 2") root = tracking.AutoCheckpointable() root.asset1 = tracking.TrackableAsset(file1) root.asset2 = tracking.TrackableAsset(file2) save_dir = os.path.join(self.get_temp_dir(), "save_dir") save.save(root, save_dir, signatures={}) file_io.delete_file(file1) file_io.delete_file(file2) load_dir = os.path.join(self.get_temp_dir(), "load_dir") file_io.rename(save_dir, load_dir) imported = load.load(load_dir) with open(imported.asset1.asset_path.numpy(), "r") as f: self.assertEquals("contents 1", f.read()) with open(imported.asset2.asset_path.numpy(), "r") as f: self.assertEquals("contents 2", f.read())
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 testTrtGraphConverter_BasicConversion_v2(self): """Test case for trt_convert.TrtGraphConverter().""" if not is_tensorrt_enabled(): return np_input = np.random.random_sample([4, 1, 1]).astype(np.float32) # Create a model and save it. input_saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir()) root = self._GetModelForV2() expected_output = root.run(np_input) save.save(root, input_saved_model_dir, {_SAVED_MODEL_SIGNATURE_KEY: root.run}) # Run TRT conversion. converter = trt_convert.TrtGraphConverterV2( input_saved_model_dir=input_saved_model_dir, input_saved_model_signature_key=_SAVED_MODEL_SIGNATURE_KEY, conversion_params=trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace( precision_mode=trt_convert.TrtPrecisionMode.FP32, is_dynamic_op=True, maximum_cached_engines=2, use_function_backup=False)) converted_func = converter.convert() def _check_trt_ops(graph_def): trt_op_names = [ node.name for node in graph_def.node if node.op == "TRTEngineOp" ] for func in graph_def.library.function: for node in func.node_def: if node.op == "TRTEngineOp": trt_op_names.append(node.name) self.assertEqual(1, len(trt_op_names)) self.assertIn("TRTEngineOp_0", trt_op_names[0]) # Verify the converted GraphDef and ConcreteFunction. self.assertIsInstance(converted_func, def_function.Function) converted_concrete_func = converted_func.get_concrete_function( tensor_spec.TensorSpec(shape=[None, 1, 1], dtype=dtypes.float32)) _check_trt_ops(converted_concrete_func.graph.as_graph_def()) # Save the converted model without any TRT engine cache. output_saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir()) converter.save(output_saved_model_dir) unexpected_asset_file = os.path.join( output_saved_model_dir, "assets/trt-serialized-engine.TRTEngineOp_0") self.assertFalse(os.path.exists(unexpected_asset_file)) # Run the converted function to populate the engine cache. output_with_trt = converted_func(np_input) self.assertEqual(1, len(output_with_trt)) self.assertAllClose( expected_output, output_with_trt[0], atol=1e-6, rtol=1e-6) # Save the converted model again with serialized engine cache. output_saved_model_dir = tempfile.mkdtemp(dir=self.get_temp_dir()) converter.save(output_saved_model_dir) expected_asset_file = os.path.join( output_saved_model_dir, "assets/trt-serialized-engine.TRTEngineOp_0") self.assertTrue(os.path.exists(expected_asset_file)) self.assertTrue(os.path.getsize(expected_asset_file)) # Load and verify the converted model. # # TODO(laigd): the name of then new input_signature of the # `root_with_trt.run` function is empty string (originaly was None), # investigate why. root_with_trt = load.load(output_saved_model_dir) # TODO(laigd): `root_with_trt.run` is still using the original graph without # trt. Consider changing that. # _check_trt_ops( # root_with_trt.run.get_concrete_function().graph.as_graph_def()) converted_signature = root_with_trt.signatures[_SAVED_MODEL_SIGNATURE_KEY] _check_trt_ops(converted_signature.graph.as_graph_def()) output_with_trt = converted_signature(ops.convert_to_tensor(np_input)) # The output of running the converted signature is a dict due to # compatibility reasons with V1 SavedModel signature mechanism. output_with_trt = output_with_trt[output_with_trt.keys()[0]] self.assertAllClose(expected_output, output_with_trt, atol=1e-6, rtol=1e-6)
def func(x): if not hasattr(closure, "model"): closure.model = load.load(path) return closure.model.f(x)
def test_restore_output_shapes(self): saved = self._v1_output_shape_saved_model() imported = load.load(saved) fn = imported.signatures["serving_default"] self.assertEqual(tensor_shape.TensorShape([1]), fn.outputs[0].shape)
def test_version_info(self): path = self._signature_with_no_inputs() imported = load.load(path) self.assertEqual(versions.__version__, imported.tensorflow_version) self.assertEqual(versions.__git_version__, imported.tensorflow_git_version)
def test_unfed_placeholder_exception(self): path = self._unfed_placeholder_signature() with self.assertRaisesRegexp( lift_to_graph.UnliftableError, "signature needs an input for each placeholder.*\n\nUnable to lift"): load.load(path)
def test_custom_pruning(self): path = self._no_signatures_model() root = load.load(path) fn = root.prune("x:0", "out:0") self.assertEqual(2., self.evaluate(fn(x=array_ops.ones([])))) root.graph.as_graph_element("x:0")
def cycle(self, obj, signatures=None): path = tempfile.mkdtemp(prefix=self.get_temp_dir()) save.save(obj, path, signatures=signatures or {}) return load.load(path)
def test_ref_variable_import(self): saved = self._v1_single_metagraph_saved_model(use_resource=False) imported = load.load(saved) fn = imported.signatures["serving_default"] self.assertEqual(6., fn(start=constant_op.constant(2.))["output"].numpy())
def test_cond(self): first_path = self._v1_cond_saved_model() imported = load.load(first_path) function = imported.signatures["serving_default"] self.assertAllClose({"output": 1.}, function(constant_op.constant(True))) self.assertAllClose({"output": 0.}, function(constant_op.constant(False)))
def test_nested_while(self): first_path = self._v1_nested_while_saved_model() imported = load.load(first_path) function = imported.signatures["serving_default"] self.assertAllClose({"output": 20}, function(constant_op.constant(4))) self.assertAllClose({"output": 35}, function(constant_op.constant(5)))
def test_no_signature(self): path = self._no_signatures_model() imported = load.load(path) self.assertEqual([], list(imported.signatures.keys()))
def test_signature_with_no_inputs(self): path = self._signature_with_no_inputs() imported = load.load(path) self.assertEqual([2], imported.signatures["key"]()["value"].shape)