[Python] MobileNet dynamic range quantization breaks model accuracy

I've been working on building a binary fire detection model that will run on ESP32s based on TF keras MobileNet v2 and transfer trained from ImageNet weights. After training the model and inferencing it works well with about 95% accuracy, but after converting to TF Lite format and performing dynamic range quantization it always classifies as a single label.

train_dataset = tf.keras.utils.image_dataset_from_directory(train_dir,
seed=123,
image_size=IMG_SIZE,
batch_size=BATCH_SIZE)

validation_dataset = tf.keras.utils.image_dataset_from_directory(test_dir,
seed=123,
image_size=IMG_SIZE,
batch_size=BATCH_SIZE)

class_names = train_dataset.class_names

val_batches = tf.data.experimental.cardinality(validation_dataset)
test_dataset = validation_dataset.take(val_batches // 5)
validation_dataset = validation_dataset.skip(val_batches // 5)

AUTOTUNE = tf.data.AUTOTUNE

train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
validation_dataset = validation_dataset.prefetch(buffer_size=AUTOTUNE)
test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

data_augmentation = tf.keras.Sequential([
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.2),
])

preprocess_input = tf.keras.applications.mobilenet_v2.preprocess_input

# Create the base model from the pre-trained model MobileNet V2
IMG_SHAPE = IMG_SIZE + (3,)
base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
include_top=False,
alpha=ALPHA,
weights='imagenet')
#depth_multiplier=RESMUL,

image_batch, label_batch = next(iter(train_dataset))
feature_batch = base_model(image_batch)

base_model.trainable = True

# Let's take a look to see how many layers are in the base model
print("Number of layers in the base model: ", len(base_model.layers))

fine_tune_at = 100

for layer in base_model.layers[:fine_tune_at]:
layer.trainable = False

global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
feature_batch_average = global_average_layer(feature_batch)
print(feature_batch_average.shape)

prediction_layer = tf.keras.layers.Dense(1, activation='sigmoid')
prediction_batch = prediction_layer(feature_batch_average)
print(prediction_batch.shape)

inputs = tf.keras.Input(shape=IMG_SHAPE)
x = data_augmentation(inputs)
x = preprocess_input(x)
x = base_model(x, training=False)
x = global_average_layer(x)
x = tf.keras.layers.Dropout(0.2)(x)
outputs = prediction_layer(x)
model = tf.keras.Model(inputs, outputs)

base_learning_rate = 0.0001

model.compile(loss=tf.keras.losses.BinaryCrossentropy(),
optimizer = tf.keras.optimizers.RMSprop(learning_rate=base_learning_rate/10),
metrics=[tf.keras.metrics.BinaryAccuracy(threshold=0.5, name='accuracy')])

loss0, accuracy0 = model.evaluate(validation_dataset)


Here is the code for quantization:

def convert_to_tflite_with_calibration(keras_model_path, output_path, calibration_dataset):
model = tf.keras.models.load_model(keras_model_path)
converter = tf.lite.TFLiteConverter.from_keras_model(model)

def representative_dataset():
for images, _ in calibration_dataset.take(100): # only a subset for calibration
yield [images]

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset

# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

tflite_model = converter.convert()
with open(output_path, 'wb') as f:
f.write(tflite_model)

return tflite_model

# Define paths for different model variants
base_tflite_path = 'firedetection.tflite'
quant_tflite_path = 'firedetection_quantized.tflite'


# Convert to TFLite with dynamic range quantization
convert_to_tflite_with_calibration(keras_model_path, quant_tflite_path, quantize=True)


Any ideas on what I'm doing wrong?

Post-training + TF Lite conversion

Tested on 100 images
Overall accuracy: 97.00%


Post-training + TF Lite conversion + Dynamic range quantization

Per-class accuracy:
neutral: 0.0000 (0/3191)
fire: 1.0000 (1115/1115)

Overall accuracy: 0.2589%

Continue reading...