def test_compile_simple_fp16(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) compiled_1 = compiler.compile_source( source=SavedModelFile(model_path=model_dir), config=Config(input_formats=[])) compiled_2 = compiler.compile_source( source=SavedModelFile(model_path=model_dir), config=Config(optimization=True, supported_types=[tf.float16], input_formats=[''])) self.assertLess(len(compiled_1.tflite_model), len(compiled_2.tflite_model))
def test_compile_with_signature_keys(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) config = Config(signature_keys='serving_default', max_batch_size=4, input_formats=['channels_last', 'channels_last']) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NHWC, dims=[4]), # pylint: disable=no-member ModelInput(name='y', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NHWC, dims=[4]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=TfDataType.DT_FLOAT, dims=[4]) ])
def test_compile_int8_with_optimize_offline(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) optimization_input_fn = _input_fn calibration_input_fn = _input_fn config = Config(max_batch_size=4, input_formats=['channels_last', 'channels_last'], enable_int8=True, calibration_input_fn=calibration_input_fn, optimize_offline=True, input_fn=optimization_input_fn) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NHWC, dims=[4]), # pylint: disable=no-member ModelInput(name='y', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NHWC, dims=[4]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=TfDataType.DT_FLOAT, dims=[4]) ])
def test_compile_with_two_saved_model_tags(self): with TemporaryDirectory() as model_dir: _save_saved_model_file_with_two_tags(model_dir) config = Config.from_json({ 'input_names': ['x', 'y'], 'output_names': ['z'], 'enable_nhwc_to_nchw': True, 'max_batch_size': 1, 'saved_model_tags': ['serve', 'graph2'] }) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=tf.float32.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[4, 2, 3]), # pylint: disable=no-member ModelInput(name='y', data_type=tf.float32.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[4, 2, 3]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=tf.float32.as_datatype_enum, dims=[4, 2, 3]) ])
def test_compile_with_all_params_with_shape(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) config = Config.from_json({ 'input_names': ['x', 'y'], 'input_shapes': [[1, 2, 3, 4], [1, 2, 3, 4]], 'output_names': ['z'], 'enable_nhwc_to_nchw': False }) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=tf.float32.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[2, 3, 4]), # pylint: disable=no-member ModelInput(name='y', data_type=tf.float32.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[2, 3, 4]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=tf.float32.as_datatype_enum, dims=[2, 3, 4]) ])
def test_compile_with_fp16(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) config = Config.from_json({ 'max_batch_size': 1, 'data_type': 'FP16' }) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=tf.float16.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[2, 3, 4]), # pylint: disable=no-member ModelInput(name='y', data_type=tf.float16.as_datatype_enum, format=ModelInput.FORMAT_NONE, dims=[2, 3, 4]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=tf.float16.as_datatype_enum, dims=[2, 3, 4]) ])
def test_compile_select_tf_op(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) compiled = compiler.compile_source( source=SavedModelFile(model_path=model_dir), config=Config(supported_ops=[tf.lite.OpsSet.SELECT_TF_OPS], input_formats=['channels_first'])) self.assertIsInstance(compiled.tflite_model, bytes)
def test_compile_simple(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) compiled = compiler.compile_source( source=SavedModelFile(model_path=model_dir), config=Config( input_formats=['channels_first', 'channels_last'])) self.assertIsInstance(compiled.tflite_model, bytes)
def test_compile_simple(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) compiled = compiler.compile_source(source=SavedModelFile(model_path=model_dir), config=Config.from_json({"input_formats": ['channels_first']})) self.assertEqual(len(compiled.inputs), 2) self.assertEqual([model_input.data_format for model_input in compiled.inputs], [DataFormat.CHANNELS_FIRST, None]) self.assertEqual([model_output.name for model_output in compiled.outputs], ['z'])
def test_compile_with_less_input_formats(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) config = Config(max_batch_size=4, input_formats=['channels_last']) with self.assertRaises(ValueError) as context_manager: compiler.compile_source(SavedModelFile(model_path=model_dir), config) self.assertEqual( context_manager.exception.args, ('Number of input formats (1) does not match number of inputs (2)', ))
def test_compile_with_signature_keys(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) compiled = compiler.compile_source(source=SavedModelFile(model_path=model_dir), config=Config.from_json({ "input_names": ['x', 'y'], "output_names": ['z'], "input_formats": ['channels_first'], "signature_keys": 'serving_default'})) self.assertEqual([model_input.tensor.name for model_input in compiled.inputs], ['x', 'y']) self.assertEqual([model_input.data_format for model_input in compiled.inputs], [DataFormat.CHANNELS_FIRST, None]) self.assertEqual([model_output.name for model_output in compiled.outputs], ['z'])
def _make_tftrt_saved_model(model_dir) -> TfTRTSavedModel: with tf.compat.v1.Session(graph=tf.Graph()) as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 4], name='y') weight = tf.Variable(initial_value=[2.0, 3.0, 4.0, 5.0], dtype=tf.float32) output_z = tf.add(input_x + input_y, weight, name='z') session.run(weight.initializer) builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(model_dir) input_x_tensor_info = tf.compat.v1.saved_model.build_tensor_info( input_x) input_y_tensor_info = tf.compat.v1.saved_model.build_tensor_info( input_y) output_tensor_info = tf.compat.v1.saved_model.build_tensor_info( output_z) builder.add_meta_graph_and_variables( sess=session, tags=[tf.compat.v1.saved_model.tag_constants.SERVING], signature_def_map={ 'serving_default': tf.compat.v1.saved_model.build_signature_def( inputs={ 'x': input_x_tensor_info, 'y': input_y_tensor_info }, outputs={'z': output_tensor_info}) }, clear_devices=True) builder.save() compiler, config_type = compiler_repository.REPOSITORY.get( SavedModelFile, TfTRTSavedModel) return compiler( source=SavedModelFile(model_path=model_dir), # pylint: disable=too-many-function-args config=config_type.from_json({ 'max_batch_size': 4, # pylint: disable=too-many-function-args 'input_formats': ['channels_last', 'channels_last'] }))
def test_compile_with_input_formats_is_none(self): with TemporaryDirectory() as model_dir: _save_saved_model_file(model_dir) config = Config(max_batch_size=4, input_formats=None) compiled = compiler.compile_source( SavedModelFile(model_path=model_dir), config) self.assertEqual( compiled.get_inputs(), [ ModelInput(name='x', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NONE, dims=[4]), # pylint: disable=no-member ModelInput(name='y', data_type=TfDataType.DT_FLOAT, format=ModelInput.FORMAT_NONE, dims=[4]) ]) # pylint: disable=no-member self.assertEqual(compiled.get_outputs(), [ ModelOutput(name='z', data_type=TfDataType.DT_FLOAT, dims=[4]) ])
def test_from_json(self): config = SavedModelFile.from_json({'input_model': 'foo'}) self.assertEqual(config.model_path, 'foo')
def test_from_env(self): config = SavedModelFile.from_env({'SAVED_MODEL_PATH': 'model'}) self.assertEqual(config.model_path, 'model')