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')
