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)
