def test_export_tflite(self): input_dim = 4 self.model.model = test_util.build_model(input_shape=[input_dim], num_classes=2) tflite_file = os.path.join(self.get_temp_dir(), 'model.tflite') self.model._export_tflite(tflite_file) self._test_tflite(self.model.model, tflite_file, input_dim)
def test_export_tflite_quantized(self): input_dim = 4000 num_classes = 2 max_input_value = 5 model = test_util.build_model([input_dim], num_classes) tflite_file = os.path.join(self.get_temp_dir(), 'model_quantized.tflite') dataloader = test_util.get_dataloader( data_size=1, input_shape=[input_dim], num_classes=num_classes, max_input_value=max_input_value) representative_dataset = dataloader.gen_dataset( batch_size=1, is_training=False) quantization_types = (QuantizationType.DYNAMIC, QuantizationType.INT8, QuantizationType.FP16, QuantizationType.FP32) model_sizes = (9088, 9600, 17280, 32840) for quantization_type, model_size in zip(quantization_types, model_sizes): model_util.export_tflite( model, tflite_file, quantization_type=quantization_type, representative_dataset=representative_dataset) self._test_tflite( model, tflite_file, input_dim, max_input_value, atol=1e-01) self.assertNear(os.path.getsize(tflite_file), model_size, 300) quant_configs = _get_quantization_config_list(input_dim, num_classes, max_input_value) model_sizes = [9088, 9536, 9600, 17280] for config, model_size in zip(quant_configs, model_sizes): model_util.export_tflite(model, tflite_file, quantization_config=config) self._test_tflite( model, tflite_file, input_dim, max_input_value, atol=1e-01) self.assertNear(os.path.getsize(tflite_file), model_size, 300)
def test_predict_top_k(self): input_shape = [24, 24, 3] self.model.model = test_util.build_model(input_shape, self.num_classes) ds = test_util.get_dataloader(2, input_shape, self.num_classes) raw_data = tf.random.uniform(shape=[2] + input_shape, minval=0, maxval=1, dtype=tf.float32) for data in [ds, raw_data]: topk_results = self.model.predict_top_k(data, k=2, batch_size=1) for topk_result in topk_results: top1_result, top2_result = topk_result[0], topk_result[1] top1_label, top1_prob = top1_result[0], top1_result[1] top2_label, top2_prob = top2_result[0], top2_result[1] self.assertIn(top1_label, self.model.index_to_label) self.assertIn(top2_label, self.model.index_to_label) self.assertNotEqual(top1_label, top2_label) self.assertLessEqual(top1_prob, 1) self.assertGreaterEqual(top1_prob, top2_prob) self.assertGreaterEqual(top2_prob, 0) self.assertEqual(top1_prob + top2_prob, 1.0)
def test_export_saved_model(self): self.model.model = test_util.build_model(input_shape=[4], num_classes=2) saved_model_filepath = os.path.join(self.get_temp_dir(), 'saved_model/') self.model._export_saved_model(saved_model_filepath) self.assertTrue(os.path.isdir(saved_model_filepath)) self.assertNotEqual(len(os.listdir(saved_model_filepath)), 0)
def test_export_tfjs(self): input_dim = 4000 num_classes = 2 model = test_util.build_model([input_dim], num_classes) output_dir = os.path.join(self.get_temp_dir(), 'tfjs') model_util.export_tfjs(model, output_dir) self.assertTrue(os.path.exists(output_dir)) expected_model_json = os.path.join(output_dir, 'model.json') self.assertTrue(os.path.exists(expected_model_json))
def test_export_tflite_quantized(self): input_dim = 4 num_classes = 2 max_input_value = 5 model = test_util.build_model([input_dim], num_classes) tflite_file = os.path.join(self.get_temp_dir(), 'model_quantized.tflite') for config in _get_quantization_config_list(input_dim, num_classes, max_input_value): model_util.export_tflite(model, tflite_file, config, _mock_gen_dataset) self._test_tflite( model, tflite_file, input_dim, max_input_value, atol=1e-01)
def test_export_tfjs_saved_model(self): input_dim = 4000 num_classes = 2 model = test_util.build_model([input_dim], num_classes) saved_model_dir = os.path.join(self.get_temp_dir(), 'saved_model_for_js') model.save(saved_model_dir) output_dir = os.path.join(self.get_temp_dir(), 'tfjs') model_util.export_tfjs(saved_model_dir, output_dir) if model_util.HAS_TFJS: self.assertTrue(os.path.exists(output_dir)) expected_model_json = os.path.join(output_dir, 'model.json') self.assertTrue(os.path.exists(expected_model_json))
def test_export_tflite_quantized(self): input_dim = 4000 num_classes = 2 max_input_value = 5 model = test_util.build_model([input_dim], num_classes) tflite_file = os.path.join(self.get_temp_dir(), 'model_quantized.tflite') quant_configs = _get_quantization_config_list(input_dim, num_classes, max_input_value) model_sizes = [9088, 9536, 9600, 17280] for config, model_size in zip(quant_configs, model_sizes): model_util.export_tflite(model, tflite_file, config, _mock_gen_dataset) self._test_tflite(model, tflite_file, input_dim, max_input_value, atol=1e-01) self.assertNear(os.path.getsize(tflite_file), model_size, 300)
def test_predict_top_k(self): input_shape = [24, 24, 3] self.model.model = test_util.build_model(input_shape, self.num_classes) data = test_util.get_dataloader(2, input_shape, self.num_classes) topk_results = self.model.predict_top_k(data, k=2, batch_size=1) for topk_result in topk_results: top1_result, top2_result = topk_result[0], topk_result[1] top1_label, top1_prob = top1_result[0], top1_result[1] top2_label, top2_prob = top2_result[0], top2_result[1] self.assertIn(top1_label, self.model.index_to_label) self.assertIn(top2_label, self.model.index_to_label) self.assertNotEqual(top1_label, top2_label) self.assertLessEqual(top1_prob, 1) self.assertGreaterEqual(top1_prob, top2_prob) self.assertGreaterEqual(top2_prob, 0) self.assertEqual(top1_prob + top2_prob, 1.0)
def test_export_tflite_quantized(self): input_dim = 4 num_classes = 2 max_input_value = 5 self.model.model = test_util.build_model([input_dim], num_classes) tflite_file = os.path.join(self.get_temp_dir(), 'model_quantized.tflite') self.model._export_tflite(tflite_file, quantized=True, quantization_steps=1, representative_data=test_util.get_dataloader( data_size=1, input_shape=[input_dim], num_classes=num_classes, max_input_value=max_input_value)) self._test_tflite(self.model.model, tflite_file, input_dim, max_input_value, atol=1e-01)
def setUp(self): super(CustomModelTest, self).setUp() self.model = MockCustomModel(model_spec=None, shuffle=False) self.model.model = test_util.build_model(input_shape=[4], num_classes=2)