def test_custom_model_to_debug(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1,),
num_last_blocks=1,
)
model = backbone.backbone
inputs = model.input / 255.0
custom_model = tf.keras.Model(inputs=model.input, outputs=model(inputs))
debug_model = tflite_debugger.convert_to_debug_model(custom_model)
self.assertEqual(len(debug_model.output_names), 18)
input_image = tf.keras.Input(
shape=(image_dim, image_dim, 3), name="input_image"
)
outputs = model(input_image)
custom_model = tf.keras.Model(inputs=input_image, outputs=outputs)
debug_model = tflite_debugger.convert_to_debug_model(custom_model)
self.assertEqual(len(debug_model.output_names), 17) # no normalization layer
def build_fpn(self, num_first_blocks: int):
backbone = resnet.ResNetBackbone(
input_shape=(self.image_dim, self.image_dim, 3),
units_per_block=(1, 1),
num_last_blocks=2,
)
return fpn.FPNBackbone(backbone,
depth=32,
num_first_blocks=num_first_blocks)
def test_model_with_box_shape(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1),
num_last_blocks=1,
)
tasks = [standard_tasks.get_box_shape_task()]
builder = FPNBuilder(backbone=backbone, tasks=tasks)
self.train_export_model(builder)
def test_quantize_resnet_forward(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1),
num_last_blocks=2,
)
inputs = tf.keras.layers.Input(shape=[image_dim, image_dim, 3])
feature_maps = backbone.forward(inputs)
quantized_feature_maps = backbone.forward(inputs, quantized=True)
def test_model_with_objectness_and_fpn_backbone(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1, 1),
num_last_blocks=3,
)
backbone = FPNBackbone(backbone, depth=16, num_first_blocks=1)
tasks = [standard_tasks.get_objectness_task()]
builder = FPNBuilder(backbone=backbone, tasks=tasks)
self.train_export_model(builder)
def test_model_with_many_tasks_many_fms(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1, 1),
num_last_blocks=2,
)
tasks = [
standard_tasks.get_objectness_task(),
standard_tasks.get_box_shape_task(),
standard_tasks.get_multiclass_task(num_classes=10),
]
builder = FPNBuilder(backbone=backbone, tasks=tasks)
self.train_export_model(builder)
def test_train_and_export_model(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(None, None, 3),
units_per_block=(1, 1),
num_last_blocks=1,
)
num_scales = 3
se_builder = bse.BoxSizeEstimatorBuilder(
(image_dim, image_dim, 3),
backbone.backbone,
box_size_task=bse.get_mean_box_size_task(),
objectness_task=bse.get_objectness_task(),
num_scales=num_scales)
self.train_export_model(se_builder)
def test_create_box_size_estimator_backbone(self):
image_dim = 128
backbone = resnet.ResNetBackbone(
input_shape=(None, None, 3),
units_per_block=(1, 1),
num_last_blocks=1,
)
se_backbone = bse.SizeEstimatorBackbone(backbone.backbone,
(image_dim, image_dim, 3),
num_scales=3)
se_backbone_model = se_backbone.as_model()
tflite_ops.TFLiteModel.from_keras_model(se_backbone_model)
se_backbone_model.summary()
se_backbone_model = se_backbone.as_model(quantized=True)
tflite_ops.TFLiteModel.from_keras_model(se_backbone_model)
def test_model_conversion(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, ),
num_last_blocks=1,
)
model = backbone.backbone
qmodel = tfmot.quantization.keras.quantize_model(model)
tflite_ops.TFLiteModel.from_keras_model(model)
tflite_ops.TFLiteModel.from_keras_model(qmodel,
[tf.lite.Optimize.DEFAULT])
tflite_ops.TFLiteModel.from_keras_model(
qmodel, [tf.lite.Optimize.OPTIMIZE_FOR_LATENCY])
def test_convert_to_tflite_debug_model(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1),
num_last_blocks=2,
)
model = backbone.backbone
qaware_model = tfmo.quantization.keras.quantize_model(model)
debug_keras_model = tflite_debugger.convert_to_debug_model(qaware_model)
debug_tflite_model = TFLiteModel.from_keras_model(debug_keras_model)
self.assertEqual(
len(debug_keras_model.output_names), len(debug_tflite_model.output_names)
)
matches = tflite_debugger.match_debug_models_output_names(
debug_keras_model, debug_tflite_model
)
self.assertEqual(len(matches), len(debug_keras_model.output_names))
inputs = np.random.rand(*[1, image_dim, image_dim, 3])
output_diffs = tflite_debugger.diff_quantiztion_outputs(
inputs, keras_model=debug_keras_model, tflite_model=debug_tflite_model
)
self.assertEqual(len(output_diffs), len(debug_tflite_model.output_names))
def representative_dataset():
while True:
inputs = np.random.rand(*[1, image_dim, image_dim, 3])
yield inputs
output_diffs = tflite_debugger.debug_models_quantization(
representative_dataset(),
keras_model=debug_keras_model,
tflite_model=debug_tflite_model,
max_samples=2,
)
self.assertEqual(len(output_diffs), len(debug_tflite_model.output_names))
self.assertEqual(len(output_diffs[0].metrics["mae"]), 2)
def test_custom_model_to_debug_quantized(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1,),
num_last_blocks=1,
)
input_image = tf.keras.Input(
shape=(image_dim, image_dim, 3), name="input_image"
)
outputs = backbone.forward(input_image, quantized=True)
custom_model = tf.keras.Model(inputs=input_image, outputs=outputs)
debug_model = tflite_debugger.convert_to_debug_model(custom_model)
debug_model.summary()
debug_tflite_model = TFLiteModel.from_keras_model(debug_model)
print(debug_tflite_model.output_names)
def test_convert_to_debug_model(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1,),
num_last_blocks=1,
)
model = backbone.backbone
qaware_model = tfmo.quantization.keras.quantize_model(model)
debug_model = tflite_debugger.convert_to_debug_model(model)
debug_qaware_model = tflite_debugger.convert_to_debug_model(qaware_model)
self.assertEqual(len(debug_model.outputs), 17)
# quantized model has one mode layer just after input layer
self.assertEqual(len(debug_qaware_model.outputs), 17)
TFLiteModel.from_keras_model(debug_qaware_model)
def test_debug_quantized_model(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1, 1),
num_last_blocks=2,
)
tasks = [standard_tasks.get_objectness_task()]
builder = FPNBuilder(backbone=backbone, tasks=tasks)
builder.build()
image_dim = builder.input_shape[1]
raw_dataset = datasets_ops.from_numpy_generator(
utils.create_fake_detection_dataset_generator(num_steps=100))
train_dataset = datasets_ops.prepare_dataset(
raw_dataset,
model_image_size=(image_dim, image_dim),
augmentation_fn=aug_fn,
num_epochs=-1,
batch_size=2,
shuffle_buffer_size=1,
prefetch_buffer_size=1,
)
prepared_train_dataset = train_dataset.map(
builder.get_build_training_targets_fn())
def representative_dataset():
for features, labels in prepared_train_dataset:
for image in features["image"]:
yield np.expand_dims(image, 0)
quantized_model = builder.build_quantized()
outputs_diffs = tflite_debugger.debug_model_quantization(
representative_dataset(), quantized_model, max_samples=1)
print(outputs_diffs)
def test_dump_metadata(self):
image_dim = 64
backbone = resnet.ResNetBackbone(
input_shape=(image_dim, image_dim, 3),
units_per_block=(1,),
num_last_blocks=1,
)
model = backbone.backbone
model_path = Path(self.test_dir) / "model.tflite"
tflite_ops.TFLiteModel.from_keras_model(model).dump(model_path)
buffer = build_metadata(
output_types=[
OutputTensorType.OBJECTNESS,
OutputTensorType.BOX_SHAPE,
OutputTensorType.CLASSES,
], author="xd",
**ModelMetadata(
"test-name",
"",
labels=[
LabelDescription("box", "box_uuid"),
LabelDescription("other_box", "other_box_uuid"),
],
).asdict,
)
dump_metadata(model_path, buffer)
metadata = load_metadata(model_path)
exp_metadata = {
"name": "test-name",
"description": '{"name": "test-name", "task": "object_detection", "labels": [{"name": "box", "uuid": "box_uuid"}, {"name": "other_box", "uuid": "other_box_uuid"}], "output_interpretation": "default", "task_params": "", "version": "default"}',
"version": "default",
"subgraph_metadata": [
{
"input_tensor_metadata": [{"name": "image"}],
"output_tensor_metadata": [
{"name": "objectness"},
{"name": "box_shape"},
{"name": "classes"},
],
}
],
"author": "xd",
}
self.assertEqual(metadata, exp_metadata)
buffer = build_metadata(
output_types=[
OutputTensorType.OBJECTNESS,
OutputTensorType.BOX_SHAPE,
], author="xd",
**ModelMetadata(
"test-name",
"",
labels=[
LabelDescription("box", "box_uuid"),
LabelDescription("other_box", "other_box_uuid"),
],
).asdict,
)
dump_metadata(model_path, buffer)
metadata = load_metadata(model_path)
exp_metadata = {
"name": "test-name",
"description": '{"name": "test-name", "task": "object_localization", "labels": [{"name": "box", "uuid": "box_uuid"}, {"name": "other_box", "uuid": "other_box_uuid"}], "output_interpretation": "default", "task_params": "", "version": "default"}',
"version": "default",
"subgraph_metadata": [
{
"input_tensor_metadata": [{"name": "image"}],
"output_tensor_metadata": [
{"name": "objectness"},
{"name": "box_shape"},
],
}
],
"author": "xd",
}
self.assertEqual(metadata, exp_metadata)