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_gen_dataset(self):
input_dim = 8
data = test_util.get_dataloader(data_size=2,
input_shape=[input_dim],
num_classes=2)
ds = data.gen_dataset()
self.assertEqual(len(ds), 2)
for (feature, label) in ds:
self.assertTrue((tf.shape(feature).numpy() == np.array([1,
8])).all())
self.assertTrue((tf.shape(label).numpy() == np.array([1])).all())
ds2 = data.gen_dataset(batch_size=2)
self.assertEqual(len(ds2), 1)
for (feature, label) in ds2:
self.assertTrue((tf.shape(feature).numpy() == np.array([2,
8])).all())
self.assertTrue((tf.shape(label).numpy() == np.array([2])).all())
ds3 = data.gen_dataset(batch_size=2, is_training=True, shuffle=True)
self.assertEqual(ds3.cardinality(), 1)
for (feature, label) in ds3.take(10):
self.assertTrue((tf.shape(feature).numpy() == np.array([2,
8])).all())
self.assertTrue((tf.shape(label).numpy() == np.array([2])).all())
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_gen_dataset(self):
input_dim = 8
data = test_util.get_dataloader(
data_size=2, input_shape=[input_dim], num_classes=2)
ds = self.model._gen_dataset(data, batch_size=1, is_training=False)
expected = list(data.dataset.as_numpy_iterator())
for i, (feature, label) in enumerate(ds):
expected_feature = [expected[i][0]]
expected_label = [expected[i][1]]
self.assertTrue((feature.numpy() == expected_feature).any())
self.assertEqual(label.numpy(), expected_label)
def _get_quantization_config_list(input_dim, num_classes, max_input_value):
# Configuration for dynamic range quantization.
config1 = configs.QuantizationConfig.for_dynamic()
representative_data = test_util.get_dataloader(
data_size=1,
input_shape=[input_dim],
num_classes=num_classes,
max_input_value=max_input_value)
# Configuration for full integer quantization with integer only.
config2 = configs.QuantizationConfig.for_int8(
representative_data=representative_data, quantization_steps=1)
# Configuration for full integer quantization with float fallback.
config3 = configs.QuantizationConfig.for_float16()
return [config1, config2, config3]
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)
