def test_mobilenet_v2_saved_model_and_keras_model(self):
input_size = DEFAULT_INPUT_SIZE
output_size = 5
head_model = tf.keras.Sequential([
layers.Dense(units=32,
input_shape=(input_size, ),
activation='relu',
kernel_regularizer=l2(0.01),
bias_regularizer=l2(0.01)),
layers.Dense(units=output_size,
kernel_regularizer=l2(0.01),
bias_regularizer=l2(0.01)),
])
head_model.compile(loss='categorical_crossentropy', optimizer='sgd')
converter = tflite_transfer_converter.TFLiteTransferConverter(
output_size, self._default_base_model,
heads.KerasModelHead(head_model), optimizers.SGD(LEARNING_RATE),
DEFAULT_BATCH_SIZE)
models = converter._convert()
parameter_shapes = [(input_size, 32), (32, ), (32, output_size),
(output_size, )]
self.assertSignatureEqual(models['initialize'], [()], parameter_shapes)
self.assertSignatureEqual(models['bottleneck'], [(1, input_size)],
[(1, input_size)])
self.assertSignatureEqual(models['inference'],
[(1, input_size)] + parameter_shapes,
[(1, output_size)])
self.assertSignatureEqual(models['optimizer'],
parameter_shapes + parameter_shapes,
parameter_shapes)
# on-device model configuration.
num_classes = 2
learning_rate = 0.001
batch_size = 5
l2_rate = 0.0001
hidden_units = 128
input_shape = model.get_layer(encoder_layer).output.shape
base = bases.SavedModelBase(tflite_model)
head = Sequential([
Flatten(input_shape=input_shape),
Dense(units=hidden_units,
activation="relu",
kernel_regularizer=l2(l2_rate)),
Dense(units=num_classes,
activation="softmax",
kernel_regularizer=l2(l2_rate)),
])
# Optimizer is ignored by the converter!
head.compile(loss="categorical_crossentropy", optimizer="adam")
converter = TFLiteTransferConverter(num_classes,
base,
heads.KerasModelHead(head),
optimizers.SGD(learning_rate),
train_batch_size=batch_size)
converter.convert_and_save(tflite_ondevice_model)