def test_compile_with_fp16(self):
with TemporaryDirectory() as directory:
model_path = os.path.join(directory, 'model.ckpt')
_save_tensorflow_model(model_path)
config = Config.from_json({
'max_batch_size': 1,
'data_type': 'FP16'
})
compiled = compiler.compile_source(
TfModelFile(model_path=model_path), 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_with_all_params_with_enable_nhwc_to_nchw_true(self):
with TemporaryDirectory() as directory:
model_path = os.path.join(directory, 'model.ckpt')
_save_tensorflow_model(model_path)
config = Config.from_json({
'input_names': ['x', 'y'],
'output_names': ['z'],
'enable_nhwc_to_nchw': True,
'max_batch_size': 1
})
compiled = compiler.compile_source(
TfModelFile(model_path=model_path), 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_variables(self):
with TemporaryDirectory() as directory:
model_path = os.path.join(directory, 'model.ckpt')
_save_tensorflow_model(model_path)
config = Config.from_json({
'input_names': ['x:0', 'y:0'],
'output_names': ['z:0'],
'input_formats': ['channels_first', '']
})
compiled = compiler.compile_source(
TfModelFile(model_path=model_path), config)
self.assertIsInstance(compiled.session, tf.compat.v1.Session)
self.assertEqual(
[model_input.tensor.name for model_input in compiled.inputs],
['x:0', 'y:0'])
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:0'])
def test_compile_with_variables(self):
with TemporaryDirectory() as directory:
model_path = os.path.join(directory, 'model.ckpt')
_save_tensorflow_model(model_path)
config = Config.from_json({
'input_names': ['x:0', 'y:0'],
'output_names': ['z:0'],
'input_formats': ['channels_last']
})
compiled = compiler.compile_source(
TfModelFile(model_path=model_path), config)
compiled_graph = compiled.model_proto.graph
initializers = {
initializer.name
for initializer in compiled_graph.initializer
}
input_name = [
input_spec.name for input_spec in compiled_graph.input
if input_spec.name not in initializers
]
self.assertEqual(input_name, ['x:0', 'y:0'])
self.assertEqual(compiled.input_data_formats,
[DataFormat.CHANNELS_LAST])
def test_compile_with_none_input_format(self):
with TemporaryDirectory() as directory:
model_path = os.path.join(directory, 'model.ckpt')
_save_tensorflow_model(model_path)
config = Config.from_json({
'input_names': ['x:0', 'y:0'],
'output_names': ['z:0'],
'input_formats': []
})
compiled = compiler.compile_source(
TfModelFile(model_path=model_path), config)
self.assertEqual(
[model_input.data_format for model_input in compiled.inputs],
[None, None])
def test_from_json(self):
config = TfModelFile.from_json({'input_model': 'foo'})
self.assertEqual(config.model_path, 'foo')
def test_from_env(self):
config = TfModelFile.from_env({'CHECKPOINT_PATH': 'model'})
self.assertEqual(config.model_path, 'model')