def test_mobilenet_v2_saved_model_and_softmax_classifier_model(self): input_size = DEFAULT_INPUT_SIZE output_size = 5 batch_size = DEFAULT_BATCH_SIZE converter = tflite_transfer_converter.TFLiteTransferConverter( output_size, self._default_base_model, heads.SoftmaxClassifierHead(batch_size, (input_size, ), output_size), optimizers.SGD(LEARNING_RATE), batch_size) models = converter._convert() parameter_shapes = [(input_size, output_size), (output_size, )] self.assertSignatureEqual(models['initialize'], [()], parameter_shapes) self.assertSignatureEqual(models['bottleneck'], [(1, input_size)], [(1, input_size)]) self.assertSignatureEqual(models['train_head'], [(batch_size, input_size), (batch_size, output_size)] + parameter_shapes, [()] + parameter_shapes) self.assertSignatureEqual(models['inference'], [(1, input_size)] + parameter_shapes, [(1, output_size)]) self.assertSignatureEqual(models['optimizer'], parameter_shapes + parameter_shapes, parameter_shapes)
def test_mobilenet_v2_base_quantized_and_softmax_classifier(self): base_model = bases.MobileNetV2Base(quantize=True) head_model = heads.SoftmaxClassifierHead(BATCH_SIZE, BOTTLENECK_SHAPE, NUM_CLASSES) optimizer = optimizers.SGD(LEARNING_RATE) model = TransferModel(self.dataset_dir, base_model, head_model, optimizer) self.assertModelAchievesAccuracy(model, 0.80)
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)
def test_mobilenet_v2_base_and_softmax_classifier_model(self): input_size = 224 output_size = 5 batch_size = DEFAULT_BATCH_SIZE base = bases.MobileNetV2Base(image_size=input_size) head = heads.SoftmaxClassifierHead(batch_size, base.bottleneck_shape(), output_size) optimizer = optimizers.SGD(LEARNING_RATE) converter = tflite_transfer_converter.TFLiteTransferConverter( output_size, base, head, optimizer, batch_size) models = converter._convert() parameter_shapes = [(7 * 7 * 1280, output_size), (output_size, )] self.assertSignatureEqual(models['initialize'], [()], parameter_shapes) self.assertSignatureEqual(models['bottleneck'], [(1, input_size, input_size, 3)], [(1, 7, 7, 1280)]) self.assertSignatureEqual(models['train_head'], [(batch_size, 7, 7, 1280), (batch_size, output_size)] + parameter_shapes, [()] + parameter_shapes) self.assertSignatureEqual(models['inference'], [(1, 7, 7, 1280)] + 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)
This is the model architecture that we will train using Flower. """ head = tf.keras.Sequential([ tf.keras.Input(shape=(32, 32, 3)), tf.keras.layers.Conv2D(6, 5, activation="relu"), tf.keras.layers.MaxPooling2D(pool_size=(2, 2)), tf.keras.layers.Conv2D(16, 5, activation="relu"), tf.keras.layers.Flatten(), tf.keras.layers.Dense(units=120, activation="relu"), tf.keras.layers.Dense(units=84, activation="relu"), tf.keras.layers.Dense(units=10, activation="softmax"), ]) head.compile(loss="categorical_crossentropy", optimizer="sgd") """Convert the model for TFLite. Using 10 classes in CIFAR10, learning rate = 1e-3 and batch size = 32 This will generate a directory called tflite_model with five tflite models. Copy them in your Android code under the assets/model directory. """ base_path = bases.saved_model_base.SavedModelBase("identity_model") converter = TFLiteTransferConverter(10, base_path, heads.KerasModelHead(head), optimizers.SGD(1e-3), train_batch_size=32) converter.convert_and_save("tflite_model")
def main(): parser = argparse.ArgumentParser( description='Combines two TF models into a transfer learning model') parser.add_argument('--train_batch_size', help='Training batch size', type=int, default=20) parser.add_argument('--num_classes', help='Number of classes for the output', type=int, default=4) # Base model configuration. base_group = parser.add_mutually_exclusive_group(required=True) base_group.add_argument('--base_mobilenetv2', help='Use MobileNetV2 as the base model', dest='base_mobilenetv2', action='store_true') base_group.add_argument( '--base_model_dir', help='Use a SavedModel under a given path as the base model', type=str) parser.add_argument('--base_quantize', help='Whether the base model should be quantized', dest='base_quantize', action='store_true') parser.set_defaults(base_quantize=False) # Head model configuration. head_group = parser.add_mutually_exclusive_group(required=True) head_group.add_argument( '--head_model_dir', help='Use a SavedModel under a given path as the head model', type=str) head_group.add_argument('--head_softmax', help='Use SoftmaxClassifier for the head model', dest='head_softmax', action='store_true') parser.add_argument( '--head_l2_reg', help='L2 regularization parameter for SoftmaxClassifier', type=float) # Optimizer configuration. parser.add_argument('--optimizer', required=True, type=str, choices=['sgd', 'adam'], help='Which optimizer should be used') parser.add_argument('--sgd_learning_rate', help='Learning rate for SGD', type=float) parser.add_argument( '--out_model_dir', help='Where the generated transfer learning model is saved', required=True, type=str) args = parser.parse_args() if args.base_mobilenetv2: base = bases.MobileNetV2Base(quantize=args.base_quantize) else: base = bases.SavedModelBase(args.base_model_dir, quantize=args.base_quantize) if args.head_model_dir: head = heads.LogitsSavedModelHead(args.head_model_dir) else: head = heads.SoftmaxClassifierHead(args.train_batch_size, base.bottleneck_shape(), args.num_classes, l2_reg=args.head_l2_reg) if args.optimizer == 'sgd': if args.sgd_learning_rate is not None: optimizer = optimizers.SGD(args.sgd_learning_rate) else: raise RuntimeError( '--sgd_learning_rate is required when SGD is used as an optimizer' ) elif args.optimizer == 'adam': optimizer = optimizers.Adam() converter = tflite_transfer_converter.TFLiteTransferConverter( args.num_classes, base, head, optimizer, args.train_batch_size) converter.convert_and_save(args.out_model_dir)