Hello all,
I was able to successfully train a MobileNet v2 (object detection mode), quantize it, and convert it into .tflite format.
Basically, my custom model had to meet 3 criteria:
- The model is not too large (models should probably be below 17 KB)
- The model's input shape is [1 300 300 3] and the input data type is uint8 or numpy.uint8
- The model's has four outputs in THIS ORDER:
- output one shape: [1 10 4]
- output two shape: [1 10]
- output three shape: [1 10]
- output four shape: [1]
all of these outputs in float32 or numpy.float32 format
From what I can best understand, the model architecture doesn't have to be an exact match, so long as the inputs and outputs are compatible.
Here is my code:
import tensorflow as tf
from tensorflow.keras import layers, models
Load the MobileNetV2 feature vector model directly from TensorFlow
base_model = tf.keras.applications.MobileNetV2(
input_shape=(300, 300, 3), # Use 300x300 input shape as required
include_top=False,
weights='imagenet')
Freezing the base model
base_model.trainable = False
Adjust input shape to 300x300x3 and use uint8 data type
inputs = tf.keras.Input(shape=(300, 300, 3), dtype='uint8')
Use Lambda layer to cast inputs to float32
x = layers.Lambda(lambda image: tf.cast(image, tf.float32))(inputs)
Pass the cast inputs through the base model
x = base_model(x)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(1280, activation='relu')(x)
Bounding box output (10 detections, 4 coordinates each)
bbox_outputs = layers.Dense(40, activation='sigmoid')(x)
bbox_outputs = layers.Lambda(lambda t: tf.reshape(t, [1, 10, 4]), name="bbox_outputs")(bbox_outputs)
Class ID output (10 detections)
class_outputs = layers.Dense(10, activation='softmax')(x)
class_outputs = layers.Lambda(lambda t: tf.reshape(t, [1, 10]), name="class_outputs")(class_outputs)
Confidence score output (10 detections)
confidence_outputs = layers.Dense(10, activation='sigmoid')(x)
confidence_outputs = layers.Lambda(lambda t: tf.reshape(t, [1, 10]), name="confidence_outputs")(confidence_outputs)
Number of detections (single value)
num_detections = layers.Lambda(lambda t: tf.constant([10], dtype=tf.float32))(x)
num_detections = layers.Lambda(lambda t: tf.reshape(t, [1]), name="num_detections")(num_detections)
Define the outputs explicitly in the order you want:
1. Bounding boxes
2. Class IDs
3. Confidence scores
4. Number of detections
model = tf.keras.Model(inputs, [bbox_outputs, class_outputs, confidence_outputs, num_detections])
Compile the model
model.compile(optimizer='adam', loss='mean_squared_error')
Define a ConcreteFunction for the model with explicit output signatures
@tf.function(input_signature=[tf.TensorSpec([1, 300, 300, 3], tf.uint8)])
def model_signature(input_tensor):
outputs = model(input_tensor)
return {
'bbox_outputs': outputs[0],
'class_outputs': outputs[1],
'confidence_outputs': outputs[2],
'num_detections': outputs[3]
}
Convert the model to TensorFlow Lite using signatures
converter = tf.lite.TFLiteConverter.from_concrete_functions([model_signature.get_concrete_function()])
Apply float16 quantization
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16] # Use float16 quantization
Convert the model
tflite_model = converter.convert()
Save the TensorFlow Lite model
with open('/content/mobilenet_v2_custom_quantized.tflite', 'wb') as f:
f.write(tflite_model)
Load the TFLite model and check the input/output details to confirm correct mapping
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
Get the input and output details to verify correct input/output structure
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print("Input Details:", input_details)
print("Output Details:", output_details)