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