def test_from_csv(self):
model_spec = MockDetectorModelSpec('efficientdet-lite0')
csv_file = test_util.get_test_data_path('salads_ml_use.csv')
image_dir = test_util.get_test_data_path('salad_images')
train_data, validation_data, test_data = \
object_detector_dataloader.DataLoader.from_csv(csv_file, image_dir)
label_map = {1: 'Baked Goods', 2: 'Cheese', 3: 'Salad'}
# Checks the training data.
self.assertIsInstance(train_data, object_detector_dataloader.DataLoader)
self.assertLen(train_data, 1)
self.assertEqual(train_data.label_map, label_map)
self.assertTrue(os.path.isfile(train_data.annotations_json_file))
self.assertGreater(os.path.getsize(train_data.annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path('train_annotations.json')
self.assertTrue(
filecmp.cmp(train_data.annotations_json_file, expected_json_file))
dataset = train_data.gen_dataset(model_spec, batch_size=1, is_training=True)
for images, labels in dataset.take(10):
images_shape = tf.shape(images).numpy()
expected_shape = np.array([1, *model_spec.config.image_size, 3])
self.assertTrue((images_shape == expected_shape).all())
self.assertLen(labels, 15)
# Checks the validation data.
self.assertIsNone(validation_data)
# Checks the test data.
self.assertIsInstance(test_data, object_detector_dataloader.DataLoader)
self.assertLen(test_data, 2)
self.assertEqual(test_data.label_map, label_map)
self.assertTrue(os.path.isfile(test_data.annotations_json_file))
self.assertGreater(os.path.getsize(test_data.annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path('test_annotations.json')
self.assertTrue(
filecmp.cmp(test_data.annotations_json_file, expected_json_file))
dataset = test_data.gen_dataset(model_spec, batch_size=1, is_training=False)
for i, (images, labels) in enumerate(dataset):
self.assertLess(i, 2)
images_shape = tf.shape(images).numpy()
expected_shape = np.array([1, *model_spec.config.image_size, 3])
self.assertTrue((images_shape == expected_shape).all())
self.assertLen(labels, 15)
def _test_export_to_tflite_with_metadata(self, model):
model_name = 'model_with_metadata'
tflite_output_file = os.path.join(self.get_temp_dir(),
'%s.tflite' % model_name)
json_output_file = os.path.join(self.get_temp_dir(), '%s.json' % model_name)
labels_output_file = os.path.join(self.get_temp_dir(), 'label.txt')
model.export(
tflite_output_file,
labels_output_file,
with_metadata=True,
export_metadata_json_file=True)
self.assertTrue(os.path.isfile(tflite_output_file))
self.assertGreater(os.path.getsize(tflite_output_file), 0)
labels = self._load_labels(labels_output_file)
self.assertEqual(labels, ['cyan', 'magenta', 'yellow'])
if not metadata.TFLITE_SUPPORT_TOOLS_INSTALLED:
return
expected_json_file = test_util.get_test_data_path(
'mobilenet_v2_metadata.json')
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
def testEfficientDetLite0(self):
# Gets model specification.
spec = model_spec.get('efficientdet_lite0')
# Prepare data.
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
data = object_detector_dataloader.DataLoader.from_pascal_voc(
images_dir, annotations_dir, label_map)
# Train the model.
task = object_detector.create(data, spec, batch_size=1, epochs=1)
self.assertEqual(spec.config.num_classes, 2)
# Evaluate trained model
metrics = task.evaluate(data)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
# Export the model to saved model.
output_path = os.path.join(self.get_temp_dir(), 'saved_model')
task.export(self.get_temp_dir(),
export_format=ExportFormat.SAVED_MODEL)
self.assertTrue(os.path.isdir(output_path))
self.assertNotEqual(len(os.listdir(output_path)), 0)
# Export the model to the float TFLite model.
output_path = os.path.join(self.get_temp_dir(), 'float.tflite')
task.export(self.get_temp_dir(),
tflite_filename='float.tflite',
quantization_config=None,
export_format=ExportFormat.TFLITE,
with_metadata=True,
export_metadata_json_file=True)
# Checks the sizes of the float32 TFLite model files in bytes.
model_size = 13476379
self.assertNear(os.path.getsize(output_path), model_size, 50000)
json_output_file = os.path.join(self.get_temp_dir(), 'float.json')
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
expected_json_file = test_util.get_test_data_path(
'efficientdet_lite0_metadata.json')
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
# Evaluate the TFLite model.
task.evaluate_tflite(output_path, data)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
# Tests the default quantized model.
filename = 'model_quant.tflite'
output_path = os.path.join(self.get_temp_dir(), filename)
task.export(self.get_temp_dir(),
tflite_filename=filename,
export_format=ExportFormat.TFLITE)
model_size = 4312187
err = model_size * 0.05
self.assertTrue(os.path.isfile(output_path))
self.assertNear(os.path.getsize(output_path), model_size, err)
def test_from_squad(self, test_file, is_training, version_2_with_negative,
size):
path = test_util.get_test_data_path('squad_testdata')
squad_path = os.path.join(path, test_file)
model_spec = MockQAModelSpec(self.get_temp_dir())
data = text_dataloader.QuestionAnswerDataLoader.from_squad(
squad_path,
model_spec,
is_training=is_training,
version_2_with_negative=version_2_with_negative)
self.assertIsInstance(data, text_dataloader.QuestionAnswerDataLoader)
self.assertEqual(data.size, size)
self.assertEqual(data.version_2_with_negative, version_2_with_negative)
self.assertEqual(data.squad_file, squad_path)
if is_training:
self.assertEmpty(data.features)
self.assertEmpty(data.examples)
else:
self.assertNotEmpty(data.features)
self.assertIsInstance(data.features[0], squad_lib.InputFeatures)
self.assertEqual(len(data.features), len(data.examples))
self.assertIsInstance(data.examples[0], squad_lib.SquadExample)
def _test_export_to_tflite_with_metadata(self,
model,
expected_json_file=None):
model_name = 'model_with_metadata'
tflite_output_file = os.path.join(self.get_temp_dir(),
'%s.tflite' % model_name)
json_output_file = os.path.join(self.get_temp_dir(), '%s.json' % model_name)
labels_output_file = os.path.join(self.get_temp_dir(), 'labels.txt')
model.export(
self.get_temp_dir(),
'%s.tflite' % model_name,
with_metadata=True,
export_metadata_json_file=True)
self.assertTrue(os.path.isfile(tflite_output_file))
self.assertGreater(os.path.getsize(tflite_output_file), 0)
self.assertFalse(os.path.isfile(labels_output_file))
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
if expected_json_file is not None:
expected_json_file = test_util.get_test_data_path(expected_json_file)
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
def test_from_pascal_voc(self):
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
model_spec = MockDetectorModelSpec('efficientdet-lite0')
data = object_detector_dataloader.DataLoader.from_pascal_voc(
images_dir, annotations_dir, label_map)
self.assertIsInstance(data, object_detector_dataloader.DataLoader)
self.assertLen(data, 1)
self.assertEqual(data.label_map, label_map)
self.assertTrue(os.path.isfile(data.annotations_json_file))
self.assertGreater(os.path.getsize(data.annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path('annotations.json')
self.assertTrue(
filecmp.cmp(data.annotations_json_file, expected_json_file))
ds = data.gen_dataset(model_spec, batch_size=1, is_training=False)
for i, (images, labels) in enumerate(ds):
self.assertEqual(i, 0)
images_shape = tf.shape(images).numpy()
expected_shape = np.array([1, *model_spec.config.image_size, 3])
self.assertTrue((images_shape == expected_shape).all())
self.assertLen(labels, 15)
ds1 = data.gen_dataset(model_spec, batch_size=1, is_training=True)
# Comments out this assert since it fails externally.
# self.assertEqual(ds1.cardinality(), tf.data.INFINITE_CARDINALITY)
for images, labels in ds1.take(10):
images_shape = tf.shape(images).numpy()
expected_shape = np.array([1, *model_spec.config.image_size, 3])
self.assertTrue((images_shape == expected_shape).all())
self.assertLen(labels, 15)
def setUpClass(cls):
super(EfficientDetModelSpecTest, cls).setUpClass()
hub_path = test_util.get_test_data_path('fake_effdet_lite0_hub')
cls._spec = object_detector_spec.EfficientDetModelSpec(
model_name='efficientdet-lite0',
uri=hub_path,
hparams=dict(map_freq=1))
cls.model = cls._spec.create_model()
def test_pascal_voc_cache_writer(self):
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
cache_writer = dataloader_util.PascalVocCacheFilesWriter(label_map,
images_dir,
num_shards=1)
cache_files = dataloader_util.get_cache_files(
cache_dir=self.get_temp_dir(), cache_prefix_filename='pascal')
cache_writer.write_files(cache_files, annotations_dir)
# Checks the TFRecord file.
tfrecord_files = cache_files.tfrecord_files
self.assertTrue(os.path.isfile(tfrecord_files[0]))
self.assertGreater(os.path.getsize(tfrecord_files[0]), 0)
# Checks the annotation json file.
annotations_json_file = cache_files.annotations_json_file
self.assertTrue(os.path.isfile(annotations_json_file))
self.assertGreater(os.path.getsize(annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path('annotations.json')
self.assertTrue(filecmp.cmp(annotations_json_file, expected_json_file))
# Checks the meta_data file.
meta_data_file = cache_files.meta_data_file
self.assertTrue(os.path.isfile(meta_data_file))
self.assertGreater(os.path.getsize(meta_data_file), 0)
with tf.io.gfile.GFile(meta_data_file, 'r') as f:
meta_data_dict = yaml.load(f, Loader=yaml.FullLoader)
self.assertEqual(meta_data_dict['size'], 1)
self.assertEqual(meta_data_dict['label_map'], label_map)
# Checks xml_dict from `_get_xml_dict` function.
xml_dict = next(cache_writer._get_xml_dict(annotations_dir))
expected_xml_dict = {
'filename':
'2012_12.jpg',
'folder':
'',
'object': [{
'difficult': '1',
'bndbox': {
'xmin': '64',
'ymin': '64',
'xmax': '192',
'ymax': '192',
},
'name': 'person',
'truncated': '0',
'pose': '',
}],
'size': {
'width': '256',
'height': '256',
}
}
self.assertEqual(xml_dict, expected_xml_dict)
def testEfficientDetLite0(self):
# Gets model specification.
spec = model_spec.get('efficientdet_lite0')
# Prepare data.
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
data = object_detector_dataloader.DataLoader.from_pascal_voc(
images_dir, annotations_dir, label_map)
# Train the model.
task = object_detector.create(data, spec, batch_size=1, epochs=1)
self.assertEqual(spec.config.num_classes, 2)
# Evaluate trained model
metrics = task.evaluate(data, batch_size=1)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
# Export the model to saved model.
output_path = os.path.join(self.get_temp_dir(), 'saved_model')
task.export(self.get_temp_dir(),
export_format=ExportFormat.SAVED_MODEL)
self.assertTrue(os.path.isdir(output_path))
self.assertNotEqual(len(os.listdir(output_path)), 0)
# Export the model to TFLite model.
output_path = os.path.join(self.get_temp_dir(), 'float.tflite')
task.export(self.get_temp_dir(),
tflite_filename='float.tflite',
export_format=ExportFormat.TFLITE,
with_metadata=True,
export_metadata_json_file=True)
self.assertTrue(tf.io.gfile.exists(output_path))
self.assertGreater(os.path.getsize(output_path), 0)
json_output_file = os.path.join(self.get_temp_dir(), 'float.json')
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
expected_json_file = test_util.get_test_data_path(
'efficientdet_lite0_metadata.json')
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
# Export the model to quantized TFLite model.
# TODO(b/175173304): Skips the test for stable tensorflow 2.4 for now since
# it fails. Will revert this change after TF upgrade.
if tf.__version__.startswith('2.4'):
return
output_path = os.path.join(self.get_temp_dir(),
'model_quantized.tflite')
config = configs.QuantizationConfig.create_full_integer_quantization(
data, is_integer_only=True)
task.export(self.get_temp_dir(),
tflite_filename='model_quantized.tflite',
quantization_config=config,
export_format=ExportFormat.TFLITE)
self.assertTrue(os.path.isfile(output_path))
self.assertGreater(os.path.getsize(output_path), 0)
def test_trainable_varaible(self):
path = test_util.get_test_data_path("hub_module_v1_mini_train")
layer = hub_loader.HubKerasLayerV1V2(path, trainable=True)
self.assertLen(layer.trainable_variables, 2)
self.assertLen(layer.variables, 4)
layer = hub_loader.HubKerasLayerV1V2(path, trainable=False)
self.assertEmpty(layer.trainable_variables)
self.assertLen(layer.variables, 2)
def _test_export_to_tflite(self,
model,
threshold=1.0,
atol=1e-04,
expected_json_file=None):
tflite_output_file = os.path.join(self.get_temp_dir(), 'model.tflite')
model.export(self.get_temp_dir(),
export_format=ExportFormat.TFLITE,
quantization_config=None,
export_metadata_json_file=expected_json_file is not None)
self.assertTrue(tf.io.gfile.exists(tflite_output_file))
self.assertGreater(os.path.getsize(tflite_output_file), 0)
result = model.evaluate_tflite(tflite_output_file, self.test_data)
self.assertGreaterEqual(result['accuracy'], threshold)
spec = model.model_spec
if isinstance(spec, text_spec.AverageWordVecModelSpec):
random_inputs = np.random.randint(low=0,
high=len(spec.vocab),
size=(1, spec.seq_len),
dtype=np.int32)
elif isinstance(spec, text_spec.BertClassifierModelSpec):
input_word_ids = np.random.randint(low=0,
high=len(spec.tokenizer.vocab),
size=(1, spec.seq_len),
dtype=np.int32)
input_mask = np.random.randint(low=0,
high=2,
size=(1, spec.seq_len),
dtype=np.int32)
input_type_ids = np.random.randint(low=0,
high=2,
size=(1, spec.seq_len),
dtype=np.int32)
random_inputs = (input_word_ids, input_mask, input_type_ids)
else:
raise ValueError('Unsupported model_spec type: %s' %
str(type(spec)))
self.assertTrue(
test_util.is_same_output(tflite_output_file,
model.model,
random_inputs,
spec,
atol=atol))
if expected_json_file is not None:
json_output_file = os.path.join(self.get_temp_dir(), 'model.json')
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
expected_json_file = test_util.get_test_data_path(
expected_json_file)
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
def _get_data(model_spec, version):
path = test_util.get_test_data_path('squad_testdata')
train_data_path = os.path.join(path, 'train-v%s.json' % version)
validation_data_path = os.path.join(path, 'dev-v%s.json' % version)
version_2_with_negative = version.startswith('2')
train_data = text_dataloader.QuestionAnswerDataLoader.from_squad(
train_data_path,
model_spec,
is_training=True,
version_2_with_negative=version_2_with_negative)
validation_data = text_dataloader.QuestionAnswerDataLoader.from_squad(
validation_data_path,
model_spec,
is_training=False,
version_2_with_negative=version_2_with_negative)
return train_data, validation_data
def _test_export_to_tflite(self,
model,
validation_data,
threshold=0.0,
atol=1e-04,
expected_json_file=None):
tflite_output_file = os.path.join(self.get_temp_dir(), 'model.tflite')
model.export(self.get_temp_dir(),
export_format=ExportFormat.TFLITE,
quantization_config=None,
export_metadata_json_file=expected_json_file is not None)
self.assertTrue(os.path.isfile(tflite_output_file))
self.assertGreater(os.path.getsize(tflite_output_file), 0)
metric = model.evaluate_tflite(tflite_output_file, validation_data)
self.assertGreaterEqual(metric['final_f1'], threshold)
spec = model.model_spec
input_word_ids = np.random.randint(low=0,
high=len(spec.tokenizer.vocab),
size=(1, spec.seq_len),
dtype=np.int32)
input_mask = np.random.randint(low=0,
high=2,
size=(1, spec.seq_len),
dtype=np.int32)
input_type_ids = np.random.randint(low=0,
high=2,
size=(1, spec.seq_len),
dtype=np.int32)
random_inputs = (input_word_ids, input_mask, input_type_ids)
self.assertTrue(
test_util.is_same_output(tflite_output_file,
model.model,
random_inputs,
model.model_spec,
atol=atol))
if expected_json_file is not None:
json_output_file = os.path.join(self.get_temp_dir(), 'model.json')
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
expected_json_file = test_util.get_test_data_path(
expected_json_file)
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
def _create_pascal_voc(self):
# Saves the image into images_dir.
image_file_name = '2012_12.jpg'
image_data = np.random.rand(256, 256, 3)
images_dir = os.path.join(self.get_temp_dir(), 'images')
os.mkdir(images_dir)
save_path = os.path.join(images_dir, image_file_name)
image = PIL.Image.fromarray(image_data, 'RGB')
image.save(save_path)
# Gets the annonation path.
annotations_path = test_util.get_test_data_path('2012_12.xml')
annotations_dir = os.path.dirname(annotations_path)
label_map = {
1: 'person',
2: 'notperson',
}
return images_dir, annotations_dir, label_map
def testEfficientDetLite0(self):
# Gets model specification.
hub_path = test_util.get_test_data_path('fake_effdet_lite0_hub')
spec = object_detector_spec.EfficientDetModelSpec(
model_name='efficientdet-lite0', uri=hub_path)
# Prepare data.
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
data = object_detector_dataloader.DataLoader.from_pascal_voc(
images_dir, annotations_dir, label_map)
# Train the model.
task = object_detector.create(data, spec, batch_size=1, epochs=1)
# Evaluate trained model
metrics = task.evaluate(data, batch_size=1)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
def test_trainable_varaible(self):
path = test_util.get_test_data_path("hub_module_v1_mini_train")
layer = hub_loader.HubKerasLayerV1V2(path, trainable=True)
# Checks trainable variables.
self.assertLen(layer.trainable_variables, 2)
self.assertEqual(layer.trainable_variables[0].name, "a:0")
self.assertEqual(layer.trainable_variables[1].name, "b:0")
self.assertEqual(layer.variables, layer.trainable_variables)
# Checks non-trainable variables.
self.assertEmpty(layer.non_trainable_variables)
layer = hub_loader.HubKerasLayerV1V2(path, trainable=False)
# Checks trainable variables.
self.assertEmpty(layer.trainable_variables)
# Checks non-trainable variables.
self.assertLen(layer.non_trainable_variables, 2)
self.assertEqual(layer.non_trainable_variables[0].name, "a:0")
self.assertEqual(layer.non_trainable_variables[1].name, "b:0")
self.assertEqual(layer.variables, layer.non_trainable_variables)
def setup_testdata(instance):
"""Setup testdata under download_dir, and unzip data to dataset_dir."""
if not hasattr(instance, 'test_tempdir'):
instance.test_tempdir = tempfile.mkdtemp()
instance.download_dir = os.path.join(instance.test_tempdir, 'download')
# Copy zip file and unzip.
os.makedirs(instance.download_dir, exist_ok=True)
# Use existing copy of data, if exists; otherwise, download it.
try:
path = test_util.get_test_data_path('recommendation_movielens')
zip_file = os.path.join(path, 'ml-1m.zip')
shutil.copy(zip_file, instance.download_dir)
with zipfile.ZipFile(zip_file, 'r') as zfile:
zfile.extractall(
instance.download_dir) # Will generate at 'ml-1m'.
instance.dataset_dir = os.path.join(instance.download_dir, 'ml-1m')
except ValueError:
instance.dataset_dir = recommendation_demo.download_data(
instance.download_dir)
def get_cache_dir():
return os.path.join(test_util.get_test_data_path('demo'), 'testdata')
def test_csv_cache_writer(self):
label_map = {1: 'Baked Goods', 2: 'Cheese', 3: 'Salad'}
images_dir = test_util.get_test_data_path('salad_images')
cache_writer = dataloader_util.CsvCacheFilesWriter(label_map,
images_dir,
num_shards=1)
csv_file = test_util.get_test_data_path('salads_ml_use.csv')
for set_name, size in [('TRAIN', 1), ('TEST', 2)]:
with tf.io.gfile.GFile(csv_file, 'r') as f:
lines = [
line for line in csv.reader(f)
if line[0].startswith(set_name)
]
cache_files = dataloader_util.get_cache_files(
cache_dir=self.get_temp_dir(), cache_prefix_filename='csv')
cache_writer.write_files(cache_files, lines)
# Checks the TFRecord file.
tfrecord_files = cache_files.tfrecord_files
self.assertTrue(os.path.isfile(tfrecord_files[0]))
self.assertGreater(os.path.getsize(tfrecord_files[0]), 0)
# Checks the annotation json file.
annotations_json_file = cache_files.annotations_json_file
self.assertTrue(os.path.isfile(annotations_json_file))
self.assertGreater(os.path.getsize(annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path(
set_name.lower() + '_annotations.json')
self.assertTrue(
filecmp.cmp(annotations_json_file, expected_json_file))
# Checks the meta_data file.
meta_data_file = cache_files.meta_data_file
self.assertTrue(os.path.isfile(meta_data_file))
self.assertGreater(os.path.getsize(meta_data_file), 0)
with tf.io.gfile.GFile(meta_data_file, 'r') as f:
meta_data_dict = yaml.load(f, Loader=yaml.FullLoader)
self.assertEqual(meta_data_dict['size'], size)
self.assertEqual(meta_data_dict['label_map'], label_map)
# Checks xml_dict from `_get_xml_dict` function.
xml_dict = next(cache_writer._get_xml_dict(lines))
expected_xml_dict = {
'filename':
'279324025_3e74a32a84_o.jpg',
'object': [{
'name': 'Baked Goods',
'bndbox': {
'xmin': 9.1888,
'ymin': 101.982,
'xmax': 908.0176,
'ymax': 882.8832,
'name': 'Baked Goods',
},
'difficult': 0,
'truncated': 0,
'pose': 'Unspecified',
}],
'size': {
'width': 1600,
'height': 1200,
}
}
self.assertEqual(xml_dict, expected_xml_dict)
def test_load_with_defaults(self, module_name, trainable): inputs, expected_outputs = 10., 11. # Test modules perform increment op. path = test_util.get_test_data_path(module_name) layer = hub_loader.HubKerasLayerV1V2(path, trainable=trainable) output = layer(inputs) self.assertEqual(output, expected_outputs)
def testEfficientDetLite0(self):
# Gets model specification.
spec = model_spec.get('efficientdet_lite0')
# Prepare data.
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
data = object_detector_dataloader.DataLoader.from_pascal_voc(
images_dir, annotations_dir, label_map)
# Train the model.
task = object_detector.create(data, spec, batch_size=1, epochs=1)
self.assertEqual(spec.config.num_classes, 2)
# Evaluate trained model
metrics = task.evaluate(data)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
# Export the model to saved model.
output_path = os.path.join(self.get_temp_dir(), 'saved_model')
task.export(self.get_temp_dir(),
export_format=ExportFormat.SAVED_MODEL)
self.assertTrue(os.path.isdir(output_path))
self.assertNotEqual(len(os.listdir(output_path)), 0)
# Export the model to TFLite model.
output_path = os.path.join(self.get_temp_dir(), 'float.tflite')
task.export(self.get_temp_dir(),
tflite_filename='float.tflite',
quantization_type=QuantizationType.FP32,
export_format=ExportFormat.TFLITE,
with_metadata=True,
export_metadata_json_file=True)
# Checks the sizes of the float32 TFLite model files in bytes.
model_size = 13476379
self.assertNear(os.path.getsize(output_path), model_size, 50000)
json_output_file = os.path.join(self.get_temp_dir(), 'float.json')
self.assertTrue(os.path.isfile(json_output_file))
self.assertGreater(os.path.getsize(json_output_file), 0)
expected_json_file = test_util.get_test_data_path(
'efficientdet_lite0_metadata.json')
self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
# Evaluate the TFLite model.
task.evaluate_tflite(output_path, data)
self.assertIsInstance(metrics, dict)
self.assertGreaterEqual(metrics['AP'], 0)
# Export the model to quantized TFLite model.
# TODO(b/175173304): Skips the test for stable tensorflow 2.4 for now since
# it fails. Will revert this change after TF upgrade.
if tf.__version__.startswith('2.4'):
return
# Not include QuantizationType.FP32 here since we have already tested it
# above together with metadata file test.
types = (QuantizationType.INT8, QuantizationType.FP16,
QuantizationType.DYNAMIC)
# The sizes of the TFLite model files in bytes.
model_sizes = (4439987, 6840331, 4289875)
for quantization_type, model_size in zip(types, model_sizes):
filename = quantization_type.name.lower() + '.tflite'
output_path = os.path.join(self.get_temp_dir(), filename)
task.export(self.get_temp_dir(),
quantization_type=quantization_type,
tflite_filename=filename,
export_format=ExportFormat.TFLITE)
self.assertNear(os.path.getsize(output_path), model_size, 50000)
