Exemplo n.º 1
0
    def fit(self,
            trdst,
            valdst,
            nb_epochs,
            steps_per_epoch,
            batch_size=100,
            use_wn=False):

        opt = AdamWithWeightnorm() if use_wn else optimizers.Adam()
        self.model.compile(optimizer=opt, loss='mse', metrics=[psnr_tf])

        log_dir = os.path.join(self.log_dir, self.model_name)
        callback_list = [
            callbacks.ModelCheckpoint(self.weights_path,
                                      save_best_only=False,
                                      save_weights_only=True,
                                      verbose=1),
            callbacks.LearningRateScheduler(
                lambda e: self.lr_schedule(e, nb_epochs), verbose=0),
            callbacks.TensorBoard(log_dir=log_dir,
                                  histogram_freq=1,
                                  write_graph=True)
        ]

        print('Training model : %s' % (self.model_name))

        self.model.fit(
            x=trdst.batch(batch_size).prefetch(AUTOTUNE),
            epochs=nb_epochs,
            callbacks=callback_list,
            validation_data=valdst.batch(batch_size).prefetch(AUTOTUNE),
            steps_per_epoch=steps_per_epoch,
            verbose=1)

        return self
Exemplo n.º 2
0
  def test_validate_callbacks_predefined_callbacks(self):
    supported_predefined_callbacks = [
        callbacks.TensorBoard(),
        callbacks.CSVLogger(filename='./log.csv'),
        callbacks.EarlyStopping(),
        callbacks.ModelCheckpoint(filepath='./checkpoint'),
        callbacks.TerminateOnNaN(),
        callbacks.ProgbarLogger(),
        callbacks.History(),
        callbacks.RemoteMonitor()
    ]

    distributed_training_utils.validate_callbacks(
        supported_predefined_callbacks, adam.Adam())

    unsupported_predefined_callbacks = [
        callbacks.ReduceLROnPlateau(),
        callbacks.LearningRateScheduler(schedule=lambda epoch: 0.001)
    ]

    for callback in unsupported_predefined_callbacks:
      with self.assertRaisesRegex(ValueError,
                                  'You must specify a Keras Optimizer V2'):
        distributed_training_utils.validate_callbacks([callback],
                                                      v1_adam.AdamOptimizer())
Exemplo n.º 3
0
def get_callbacks():
    stop_on_nan_callback = callbacks.TerminateOnNaN()

    def decay(epoch):
        if epoch < 7:
            return 0.001
        else:
            return 0.001 * np.exp(0.1 * (7 - epoch))

    learning_rate_scheduler = callbacks.LearningRateScheduler(decay)

    return [stop_on_nan_callback, learning_rate_scheduler]
  def callableForTestLearningRateScheduler(model, test_obj, train_ds, num_epoch,
                                           steps, strategy, saving_filepath,
                                           **kwargs):

    cbks = [
        callbacks.LearningRateScheduler(
            schedule=lambda x: 1. / (1. + x), verbose=1)
    ]

    # It is expected that with `epochs=2`, the learning rate would drop to
    # 1 / (1 + 2) = 0.5.
    model.fit(x=train_ds, epochs=2, steps_per_epoch=steps, callbacks=cbks)
    test_obj.assertAllClose(
        float(K.get_value(model.optimizer.lr)), 0.5, atol=1e-8)

    # It is expected that with `epochs=4`, the learning rate would drop to
    # 1 / (1 + 4) = 0.25.
    model.fit(x=train_ds, epochs=4, steps_per_epoch=steps, callbacks=cbks)
    test_obj.assertAllClose(
        float(K.get_value(model.optimizer.lr)), 0.25, atol=1e-8)
Exemplo n.º 5
0
    def train(self):
        self.build_and_compile_model()

        # Identify last version of trained model
        files = get_checkpoint_file_list(self.config.checkpoint_dir,
                                         self.config.model_name)

        if not files:
            model_number = self.config.model_name + "_0"
        else:
            # Resume training vs new training decision
            if self.config.resume_train:
                print("Resume training from previous checkpoint")
                try:
                    self.recover_model_from_checkpoint()
                    model_number = self.model_version_number
                except FileNotFoundError:
                    print("Model not found. Creating new model")
                    model_number = self.model_version_number
                    safe_mkdir_recursive(self.config.checkpoint_dir +
                                         model_number)
            else:
                model_number = self.config.model_name + '_' + str(
                    int(files[-1].split('_')[-1]) + 1)
                os.mkdir(self.config.checkpoint_dir + model_number)

        self.trained_model_dir = self.config.checkpoint_dir + model_number + '/'

        # Get training and validation datasets from saved files
        dataset = self.get_dataset(train=True,
                                   val_split=True,
                                   random_split=False,
                                   shuffle=True,
                                   repeat_ds=True,
                                   normalize=False)
        train_ds, validation_ds, ds_lengths = dataset

        train_steps_per_epoch = int(
            math.ceil(ds_lengths[0] / self.config.batch_size))
        val_steps_per_epoch = int(
            math.ceil((ds_lengths[1] / self.config.batch_size)))

        def lr_scheduler(epoch, lr):
            decay_rate = 0.5
            if epoch % self.config.lr_scheduler == 0 and epoch:
                return lr * decay_rate
            return lr

        keras_callbacks = [
            callbacks.EarlyStopping(patience=self.config.patience,
                                    monitor='val_loss'),
            callbacks.TensorBoard(write_images=True,
                                  log_dir=self.config.checkpoint_dir +
                                  model_number + "/keras",
                                  histogram_freq=self.config.summary_freq),
            callbacks.LearningRateScheduler(lr_scheduler, verbose=1),
            CustomModelCheckpoint(filepath=os.path.join(
                self.config.checkpoint_dir + model_number,
                self.config.model_name + "_{epoch:02d}.h5"),
                                  save_weights_only=True,
                                  verbose=1,
                                  period=self.config.save_freq,
                                  extra_epoch_number=self.last_epoch_number +
                                  1),
        ]

        # Train!
        self.trainable_model.fit(train_ds,
                                 verbose=1,
                                 epochs=self.config.max_epochs,
                                 steps_per_epoch=train_steps_per_epoch,
                                 validation_steps=val_steps_per_epoch,
                                 validation_data=validation_ds,
                                 callbacks=keras_callbacks)
Exemplo n.º 6
0
def train_model_sample(model,
                       dataset,
                       expt='',
                       test_size=.1,
                       n_epoch=10,
                       batch_size=32,
                       num_gpus=None,
                       transform=None,
                       window_size=None,
                       balance_classes=True,
                       max_class_samples=None,
                       log_dir='/data/tensorboard_logs',
                       model_dir='/data/models',
                       model_name=None,
                       focal=False,
                       gamma=0.5,
                       optimizer=SGD(lr=0.01,
                                     decay=1e-6,
                                     momentum=0.9,
                                     nesterov=True),
                       lr_sched=rate_scheduler(lr=0.01, decay=0.95),
                       rotation_range=0,
                       flip=False,
                       shear=0,
                       zoom_range=0,
                       seed=None,
                       **kwargs):
    is_channels_first = K.image_data_format() == 'channels_first'

    if model_name is None:
        todays_date = datetime.datetime.now().strftime('%Y-%m-%d')
        data_name = os.path.splitext(os.path.basename(dataset))[0]
        model_name = '{}_{}_{}'.format(todays_date, data_name, expt)
    model_path = os.path.join(model_dir, '{}.h5'.format(model_name))
    loss_path = os.path.join(model_dir, '{}.npz'.format(model_name))

    train_dict, test_dict = get_data(dataset, test_size=test_size, seed=seed)

    n_classes = model.layers[-1].output_shape[1 if is_channels_first else -1]

    # the data, shuffled and split between train and test sets
    print('X_train shape:', train_dict['X'].shape)
    print('y_train shape:', train_dict['y'].shape)
    print('X_test shape:', test_dict['X'].shape)
    print('y_test shape:', test_dict['y'].shape)
    print('Output Shape:', model.layers[-1].output_shape)
    print('Number of Classes:', n_classes)

    def loss_function(y_true, y_pred):
        if isinstance(transform, str) and transform.lower() == 'disc':
            return losses.discriminative_instance_loss(y_true, y_pred)
        if focal:
            return losses.weighted_focal_loss(y_true,
                                              y_pred,
                                              gamma=gamma,
                                              n_classes=n_classes)
        return losses.weighted_categorical_crossentropy(y_true,
                                                        y_pred,
                                                        n_classes=n_classes)

    if num_gpus is None:
        num_gpus = train_utils.count_gpus()

    if num_gpus >= 2:
        batch_size = batch_size * num_gpus
        model = train_utils.MultiGpuModel(model, num_gpus)

    print('Training on {} GPUs'.format(num_gpus))

    model.compile(loss=loss_function,
                  optimizer=optimizer,
                  metrics=['accuracy'])

    if train_dict['X'].ndim == 4:
        DataGenerator = image_generators.SampleDataGenerator
        window_size = window_size if window_size else (30, 30)
    elif train_dict['X'].ndim == 5:
        DataGenerator = image_generators.SampleMovieDataGenerator
        window_size = window_size if window_size else (30, 30, 3)
    else:
        raise ValueError('Expected `X` to have ndim 4 or 5. Got',
                         train_dict['X'].ndim)

    # this will do preprocessing and realtime data augmentation
    datagen = DataGenerator(rotation_range=rotation_range,
                            shear_range=shear,
                            zoom_range=zoom_range,
                            horizontal_flip=flip,
                            vertical_flip=flip)

    # no validation augmentation
    datagen_val = DataGenerator(rotation_range=0,
                                shear_range=0,
                                zoom_range=0,
                                horizontal_flip=0,
                                vertical_flip=0)

    train_data = datagen.flow(train_dict,
                              batch_size=batch_size,
                              transform=transform,
                              transform_kwargs=kwargs,
                              window_size=window_size,
                              balance_classes=balance_classes,
                              max_class_samples=max_class_samples)

    val_data = datagen_val.flow(test_dict,
                                batch_size=batch_size,
                                transform=transform,
                                transform_kwargs=kwargs,
                                window_size=window_size,
                                balance_classes=False,
                                max_class_samples=max_class_samples)

    # fit the model on the batches generated by datagen.flow()
    loss_history = model.fit_generator(
        train_data,
        steps_per_epoch=train_data.y.shape[0] // batch_size,
        epochs=n_epoch,
        validation_data=val_data,
        validation_steps=val_data.y.shape[0] // batch_size,
        callbacks=[
            callbacks.LearningRateScheduler(lr_sched),
            callbacks.ModelCheckpoint(model_path,
                                      monitor='val_loss',
                                      verbose=1,
                                      save_best_only=True,
                                      save_weights_only=num_gpus >= 2),
            callbacks.TensorBoard(log_dir=os.path.join(log_dir, model_name))
        ])

    np.savez(loss_path, loss_history=loss_history.history)

    return model
Exemplo n.º 7
0
def train_model_siamese_daughter(model,
                                 dataset,
                                 expt='',
                                 test_size=.1,
                                 n_epoch=100,
                                 batch_size=1,
                                 num_gpus=None,
                                 crop_dim=32,
                                 min_track_length=1,
                                 neighborhood_scale_size=10,
                                 features=None,
                                 optimizer=SGD(lr=0.01,
                                               decay=1e-6,
                                               momentum=0.9,
                                               nesterov=True),
                                 log_dir='/data/tensorboard_logs',
                                 model_dir='/data/models',
                                 model_name=None,
                                 focal=False,
                                 gamma=0.5,
                                 lr_sched=rate_scheduler(lr=0.01, decay=0.95),
                                 rotation_range=0,
                                 flip=True,
                                 shear=0,
                                 zoom_range=0,
                                 seed=None,
                                 **kwargs):
    is_channels_first = K.image_data_format() == 'channels_first'

    if model_name is None:
        todays_date = datetime.datetime.now().strftime('%Y-%m-%d')
        data_name = os.path.splitext(os.path.basename(dataset))[0]
        model_name = '{}_{}_[{}]_neighs={}_epochs={}_seed={}_{}'.format(
            todays_date, data_name, ','.join(f[0] for f in sorted(features)),
            neighborhood_scale_size, n_epoch, seed, expt)
    model_path = os.path.join(model_dir, '{}.h5'.format(model_name))
    loss_path = os.path.join(model_dir, '{}.npz'.format(model_name))

    print('training on dataset:', dataset)
    print('saving model at:', model_path)
    print('saving loss at:', loss_path)

    train_dict, val_dict = get_data(dataset,
                                    mode='siamese_daughters',
                                    seed=seed,
                                    test_size=test_size)

    # the data, shuffled and split between train and test sets
    print('X_train shape:', train_dict['X'].shape)
    print('y_train shape:', train_dict['y'].shape)
    print('X_test shape:', val_dict['X'].shape)
    print('y_test shape:', val_dict['y'].shape)
    print('Output Shape:', model.layers[-1].output_shape)

    n_classes = model.layers[-1].output_shape[1 if is_channels_first else -1]

    def loss_function(y_true, y_pred):
        if focal:
            return losses.weighted_focal_loss(y_true,
                                              y_pred,
                                              gamma=gamma,
                                              n_classes=n_classes,
                                              from_logits=False)
        return losses.weighted_categorical_crossentropy(y_true,
                                                        y_pred,
                                                        n_classes=n_classes,
                                                        from_logits=False)

    if num_gpus is None:
        num_gpus = train_utils.count_gpus()

    if num_gpus >= 2:
        batch_size = batch_size * num_gpus
        model = train_utils.MultiGpuModel(model, num_gpus)

    print('Training on {} GPUs'.format(num_gpus))

    model.compile(loss=loss_function,
                  optimizer=optimizer,
                  metrics=['accuracy'])

    print('Using real-time data augmentation.')

    # this will do preprocessing and realtime data augmentation
    datagen = image_generators.SiameseDataGenerator(
        rotation_range=rotation_range,
        shear_range=shear,
        zoom_range=zoom_range,
        horizontal_flip=flip,
        vertical_flip=flip)

    datagen_val = image_generators.SiameseDataGenerator(rotation_range=0,
                                                        zoom_range=0,
                                                        shear_range=0,
                                                        horizontal_flip=0,
                                                        vertical_flip=0)

    total_train_pairs = tracking_utils.count_pairs(train_dict['y'],
                                                   same_probability=5.0)
    total_test_pairs = tracking_utils.count_pairs(val_dict['y'],
                                                  same_probability=5.0)

    train_data = datagen.flow(train_dict,
                              crop_dim=crop_dim,
                              batch_size=batch_size,
                              min_track_length=min_track_length,
                              neighborhood_scale_size=neighborhood_scale_size,
                              features=features)

    val_data = datagen_val.flow(
        val_dict,
        crop_dim=crop_dim,
        batch_size=batch_size,
        min_track_length=min_track_length,
        neighborhood_scale_size=neighborhood_scale_size,
        features=features)

    print('total_train_pairs:', total_train_pairs)
    print('total_test_pairs:', total_test_pairs)
    print('batch size:', batch_size)
    print('validation_steps: ', total_test_pairs // batch_size)

    # fit the model on the batches generated by datagen.flow()
    loss_history = model.fit_generator(
        train_data,
        steps_per_epoch=total_train_pairs // batch_size,
        epochs=n_epoch,
        validation_data=val_data,
        validation_steps=total_test_pairs // batch_size,
        callbacks=[
            callbacks.LearningRateScheduler(lr_sched),
            callbacks.ModelCheckpoint(model_path,
                                      monitor='val_loss',
                                      verbose=1,
                                      save_best_only=True,
                                      save_weights_only=num_gpus >= 2),
            callbacks.TensorBoard(log_dir=os.path.join(log_dir, model_name))
        ])

    model.save_weights(model_path)
    np.savez(loss_path, loss_history=loss_history.history)

    return model
Exemplo n.º 8
0
##############################
log_dir = "./K_YOLO/{}".format(datetime.today().strftime('%m-%d__%H-%M-%S') + "_" + args.name)
tensorboard = callbacks.TensorBoard(log_dir=log_dir, write_graph=True, write_grads=False, write_images=False, histogram_freq=0)
tensorboard.set_model(model)
checkpoint = callbacks.ModelCheckpoint(log_dir + "/model-{epoch:04d}.hdf5", period=1)
terminate = callbacks.TerminateOnNaN()

def step_decay(epoch):
    initial_lr = 1e-4
    drop = 0.05
    lrate = initial_lr * 1/(1 + drop * epoch)
    return lrate

def step_increase(epoch): # designed for 0-35
    initial_lr = 1e-8
    increase = 0.5
    lrate = initial_lr * (1 + increase)**epoch
    return lrate

lr_schedule = callbacks.LearningRateScheduler(step_decay)

model.fit_generator(
    generator=train_iterator,
    epochs=epochs,
    validation_data=test_iterator,
    validation_steps=len(test_items) // batch_size,
    callbacks=[lr_schedule, tensorboard], # checkpoint, tensorboard, lr_schedule,
    steps_per_epoch=len(train_items) // batch_size, #
    verbose=True)

Exemplo n.º 9
0
def train_model_retinanet(model,
                          dataset,
                          backbone,
                          expt='',
                          test_size=.1,
                          n_epoch=10,
                          batch_size=1,
                          num_gpus=None,
                          include_masks=False,
                          panoptic=False,
                          panoptic_weight=1,
                          anchor_params=None,
                          pyramid_levels=['P3', 'P4', 'P5', 'P6', 'P7'],
                          mask_size=(28, 28),
                          optimizer=SGD(lr=0.01,
                                        decay=1e-6,
                                        momentum=0.9,
                                        nesterov=True),
                          log_dir='/data/tensorboard_logs',
                          model_dir='/data/models',
                          model_name=None,
                          sigma=3.0,
                          alpha=0.25,
                          gamma=2.0,
                          score_threshold=0.01,
                          iou_threshold=0.5,
                          max_detections=100,
                          weighted_average=True,
                          lr_sched=rate_scheduler(lr=0.01, decay=0.95),
                          rotation_range=0,
                          flip=True,
                          shear=0,
                          zoom_range=0,
                          seed=None,
                          **kwargs):
    """Train a RetinaNet model from the given backbone

    Adapted from:
        https://github.com/fizyr/keras-retinanet &
        https://github.com/fizyr/keras-maskrcnn
    """

    is_channels_first = K.image_data_format() == 'channels_first'

    if model_name is None:
        todays_date = datetime.datetime.now().strftime('%Y-%m-%d')
        data_name = os.path.splitext(os.path.basename(dataset))[0]
        model_name = '{}_{}_{}'.format(todays_date, data_name, expt)

    model_path = os.path.join(model_dir, '{}.h5'.format(model_name))
    loss_path = os.path.join(model_dir, '{}.npz'.format(model_name))

    train_dict, test_dict = get_data(dataset, seed=seed, test_size=test_size)

    channel_axis = 1 if is_channels_first else -1
    n_classes = model.layers[-1].output_shape[channel_axis]

    if panoptic:
        n_semantic_classes = model.get_layer(
            name='semantic').output_shape[channel_axis]

    # the data, shuffled and split between train and test sets
    print('X_train shape:', train_dict['X'].shape)
    print('y_train shape:', train_dict['y'].shape)
    print('X_test shape:', test_dict['X'].shape)
    print('y_test shape:', test_dict['y'].shape)
    print('Output Shape:', model.layers[-1].output_shape)
    print('Number of Classes:', n_classes)

    if num_gpus is None:
        num_gpus = train_utils.count_gpus()

    if num_gpus >= 1e6:
        batch_size = batch_size * num_gpus
        model = train_utils.MultiGpuModel(model, num_gpus)

    print('Training on {} GPUs'.format(num_gpus))

    # evaluation of model is done on `retinanet_bbox`
    if include_masks:
        prediction_model = model
    else:
        prediction_model = retinanet_bbox(model,
                                          nms=True,
                                          anchor_params=anchor_params,
                                          panoptic=panoptic,
                                          class_specific_filter=False)

    retinanet_losses = losses.RetinaNetLosses(sigma=sigma,
                                              alpha=alpha,
                                              gamma=gamma,
                                              iou_threshold=iou_threshold,
                                              mask_size=mask_size)

    def semantic_loss(y_pred, y_true):
        return panoptic_weight * losses.weighted_categorical_crossentropy(
            y_pred, y_true, n_classes=n_semantic_classes)

    loss = {
        'regression': retinanet_losses.regress_loss,
        'classification': retinanet_losses.classification_loss
    }

    if include_masks:
        loss['masks'] = retinanet_losses.mask_loss

    if panoptic:
        loss['semantic'] = semantic_loss

    model.compile(loss=loss, optimizer=optimizer)

    if num_gpus >= 2:
        # Each GPU must have at least one validation example
        if test_dict['y'].shape[0] < num_gpus:
            raise ValueError('Not enough validation data for {} GPUs. '
                             'Received {} validation sample.'.format(
                                 test_dict['y'].shape[0], num_gpus))

        # When using multiple GPUs and skip_connections,
        # the training data must be evenly distributed across all GPUs
        num_train = train_dict['y'].shape[0]
        nb_samples = num_train - num_train % batch_size
        if nb_samples:
            train_dict['y'] = train_dict['y'][:nb_samples]
            train_dict['X'] = train_dict['X'][:nb_samples]

    # this will do preprocessing and realtime data augmentation
    datagen = image_generators.RetinaNetGenerator(
        # fill_mode='constant',  # for rotations
        rotation_range=rotation_range,
        shear_range=shear,
        zoom_range=zoom_range,
        horizontal_flip=flip,
        vertical_flip=flip)

    datagen_val = image_generators.RetinaNetGenerator(
        # fill_mode='constant',  # for rotations
        rotation_range=0,
        shear_range=0,
        zoom_range=0,
        horizontal_flip=0,
        vertical_flip=0)

    if 'vgg' in backbone or 'densenet' in backbone:
        compute_shapes = make_shapes_callback(model)
    else:
        compute_shapes = guess_shapes

    train_data = datagen.flow(train_dict,
                              seed=seed,
                              include_masks=include_masks,
                              panoptic=panoptic,
                              pyramid_levels=pyramid_levels,
                              anchor_params=anchor_params,
                              compute_shapes=compute_shapes,
                              batch_size=batch_size)

    val_data = datagen_val.flow(test_dict,
                                seed=seed,
                                include_masks=include_masks,
                                panoptic=panoptic,
                                pyramid_levels=pyramid_levels,
                                anchor_params=anchor_params,
                                compute_shapes=compute_shapes,
                                batch_size=batch_size)

    tensorboard_callback = callbacks.TensorBoard(
        log_dir=os.path.join(log_dir, model_name))

    # fit the model on the batches generated by datagen.flow()
    loss_history = model.fit_generator(
        train_data,
        steps_per_epoch=train_data.y.shape[0] // batch_size,
        epochs=n_epoch,
        validation_data=val_data,
        validation_steps=val_data.y.shape[0] // batch_size,
        callbacks=[
            callbacks.LearningRateScheduler(lr_sched),
            callbacks.ModelCheckpoint(model_path,
                                      monitor='val_loss',
                                      verbose=1,
                                      save_best_only=True,
                                      save_weights_only=num_gpus >= 2),
            tensorboard_callback,
            callbacks.ReduceLROnPlateau(monitor='loss',
                                        factor=0.1,
                                        patience=10,
                                        verbose=1,
                                        mode='auto',
                                        min_delta=0.0001,
                                        cooldown=0,
                                        min_lr=0),
            RedirectModel(
                Evaluate(val_data,
                         iou_threshold=iou_threshold,
                         score_threshold=score_threshold,
                         max_detections=max_detections,
                         tensorboard=tensorboard_callback,
                         weighted_average=weighted_average), prediction_model),
        ])

    model.save_weights(model_path)
    np.savez(loss_path, loss_history=loss_history.history)

    average_precisions = evaluate(
        val_data,
        prediction_model,
        iou_threshold=iou_threshold,
        score_threshold=score_threshold,
        max_detections=max_detections,
    )

    # print evaluation
    total_instances = []
    precisions = []
    for label, (average_precision,
                num_annotations) in average_precisions.items():
        print('{:.0f} instances of class'.format(num_annotations), label,
              'with average precision: {:.4f}'.format(average_precision))
        total_instances.append(num_annotations)
        precisions.append(average_precision)

    if sum(total_instances) == 0:
        print('No test instances found.')
    else:
        print(
            'mAP using the weighted average of precisions among classes: {:.4f}'
            .format(
                sum([a * b for a, b in zip(total_instances, precisions)]) /
                sum(total_instances)))
        print('mAP: {:.4f}'.format(
            sum(precisions) / sum(x > 0 for x in total_instances)))

    return model
Exemplo n.º 10
0
def train_model_retinanet(model,
                          dataset,
                          backbone,
                          expt='',
                          test_size=.1,
                          n_epoch=10,
                          batch_size=1,
                          num_gpus=None,
                          include_masks=False,
                          mask_size=(28, 28),
                          optimizer=SGD(lr=0.01,
                                        decay=1e-6,
                                        momentum=0.9,
                                        nesterov=True),
                          log_dir='/data/tensorboard_logs',
                          model_dir='/data/models',
                          model_name=None,
                          sigma=3.0,
                          alpha=0.25,
                          gamma=2.0,
                          score_threshold=0.01,
                          iou_threshold=0.5,
                          max_detections=100,
                          weighted_average=True,
                          lr_sched=rate_scheduler(lr=0.01, decay=0.95),
                          rotation_range=0,
                          flip=True,
                          shear=0,
                          zoom_range=0,
                          **kwargs):
    """Train a RetinaNet model from the given backbone

    Adapted from:
        https://github.com/fizyr/keras-retinanet &
        https://github.com/fizyr/keras-maskrcnn
    """
    is_channels_first = K.image_data_format() == 'channels_first'

    if model_name is None:
        todays_date = datetime.datetime.now().strftime('%Y-%m-%d')
        data_name = os.path.splitext(os.path.basename(dataset))[0]
        model_name = '{}_{}_{}'.format(todays_date, data_name, expt)
    model_path = os.path.join(model_dir, '{}.h5'.format(model_name))
    loss_path = os.path.join(model_dir, '{}.npz'.format(model_name))

    train_dict, test_dict = get_data(dataset, mode='conv', test_size=test_size)

    n_classes = model.layers[-1].output_shape[1 if is_channels_first else -1]
    # the data, shuffled and split between train and test sets
    print('X_train shape:', train_dict['X'].shape)
    print('y_train shape:', train_dict['y'].shape)
    print('X_test shape:', test_dict['X'].shape)
    print('y_test shape:', test_dict['y'].shape)
    print('Output Shape:', model.layers[-1].output_shape)
    print('Number of Classes:', n_classes)

    if num_gpus is None:
        num_gpus = train_utils.count_gpus()

    if num_gpus >= 1e6:
        batch_size = batch_size * num_gpus
        model = train_utils.MultiGpuModel(model, num_gpus)

    print('Training on {} GPUs'.format(num_gpus))

    def regress_loss(y_true, y_pred):
        # separate target and state
        regression = y_pred
        regression_target = y_true[..., :-1]
        anchor_state = y_true[..., -1]

        # filter out "ignore" anchors
        indices = tf.where(K.equal(anchor_state, 1))
        regression = tf.gather_nd(regression, indices)
        regression_target = tf.gather_nd(regression_target, indices)

        # compute the loss
        loss = losses.smooth_l1(regression_target, regression, sigma=sigma)

        # compute the normalizer: the number of positive anchors
        normalizer = K.maximum(1, K.shape(indices)[0])
        normalizer = K.cast(normalizer, dtype=K.floatx())

        return K.sum(loss) / normalizer

    def classification_loss(y_true, y_pred):
        # TODO: try weighted_categorical_crossentropy
        labels = y_true[..., :-1]
        # -1 for ignore, 0 for background, 1 for object
        anchor_state = y_true[..., -1]

        classification = y_pred
        # filter out "ignore" anchors
        indices = tf.where(K.not_equal(anchor_state, -1))
        labels = tf.gather_nd(labels, indices)
        classification = tf.gather_nd(classification, indices)

        # compute the loss
        loss = losses.focal(labels, classification, alpha=alpha, gamma=gamma)

        # compute the normalizer: the number of positive anchors
        normalizer = tf.where(K.equal(anchor_state, 1))
        normalizer = K.cast(K.shape(normalizer)[0], K.floatx())
        normalizer = K.maximum(K.cast_to_floatx(1.0), normalizer)

        return K.sum(loss) / normalizer

    def mask_loss(y_true, y_pred):
        def _mask(y_true, y_pred, iou_threshold=0.5, mask_size=(28, 28)):
            # split up the different predicted blobs
            boxes = y_pred[:, :, :4]
            masks = y_pred[:, :, 4:]

            # split up the different blobs
            annotations = y_true[:, :, :5]
            width = K.cast(y_true[0, 0, 5], dtype='int32')
            height = K.cast(y_true[0, 0, 6], dtype='int32')
            masks_target = y_true[:, :, 7:]

            # reshape the masks back to their original size
            masks_target = K.reshape(masks_target,
                                     (K.shape(masks_target)[0] *
                                      K.shape(masks_target)[1], height, width))
            masks = K.reshape(masks, (K.shape(masks)[0] * K.shape(masks)[1],
                                      mask_size[0], mask_size[1], -1))

            # batch size > 1 fix
            boxes = K.reshape(boxes, (-1, K.shape(boxes)[2]))
            annotations = K.reshape(annotations, (-1, K.shape(annotations)[2]))

            # compute overlap of boxes with annotations
            iou = overlap(boxes, annotations)
            argmax_overlaps_inds = K.argmax(iou, axis=1)
            max_iou = K.max(iou, axis=1)

            # filter those with IoU > 0.5
            indices = tf.where(K.greater_equal(max_iou, iou_threshold))
            boxes = tf.gather_nd(boxes, indices)
            masks = tf.gather_nd(masks, indices)
            argmax_overlaps_inds = tf.gather_nd(argmax_overlaps_inds, indices)
            argmax_overlaps_inds = K.cast(argmax_overlaps_inds, 'int32')
            labels = K.gather(annotations[:, 4], argmax_overlaps_inds)
            labels = K.cast(labels, 'int32')

            # make normalized boxes
            x1 = boxes[:, 0]
            y1 = boxes[:, 1]
            x2 = boxes[:, 2]
            y2 = boxes[:, 3]
            boxes = K.stack([
                y1 / (K.cast(height, dtype=K.floatx()) - 1),
                x1 / (K.cast(width, dtype=K.floatx()) - 1),
                (y2 - 1) / (K.cast(height, dtype=K.floatx()) - 1),
                (x2 - 1) / (K.cast(width, dtype=K.floatx()) - 1),
            ],
                            axis=1)

            # crop and resize masks_target
            # append a fake channel dimension
            masks_target = K.expand_dims(masks_target, axis=3)
            masks_target = tf.image.crop_and_resize(masks_target, boxes,
                                                    argmax_overlaps_inds,
                                                    mask_size)

            # remove fake channel dimension
            masks_target = masks_target[:, :, :, 0]

            # gather the predicted masks using the annotation label
            masks = tf.transpose(masks, (0, 3, 1, 2))
            label_indices = K.stack([tf.range(K.shape(labels)[0]), labels],
                                    axis=1)
            masks = tf.gather_nd(masks, label_indices)

            # compute mask loss
            mask_loss = K.binary_crossentropy(masks_target, masks)
            normalizer = K.shape(masks)[0] * K.shape(masks)[1] * K.shape(
                masks)[2]
            normalizer = K.maximum(K.cast(normalizer, K.floatx()), 1)
            mask_loss = K.sum(mask_loss) / normalizer

            return mask_loss

        # if there are no masks annotations, return 0; else, compute the masks loss
        return tf.cond(
            K.any(K.equal(K.shape(y_true), 0)), lambda: K.cast_to_floatx(0.0),
            lambda: _mask(y_true,
                          y_pred,
                          iou_threshold=iou_threshold,
                          mask_size=mask_size))

    # evaluation of model is done on `retinanet_bbox`
    if include_masks:
        prediction_model = model
    else:
        prediction_model = retinanet_bbox(model,
                                          nms=True,
                                          class_specific_filter=False)

    loss = {'regression': regress_loss, 'classification': classification_loss}

    if include_masks:
        loss['masks'] = mask_loss

    model.compile(loss=loss, optimizer=optimizer)

    if num_gpus >= 2:
        # Each GPU must have at least one validation example
        if test_dict['y'].shape[0] < num_gpus:
            raise ValueError('Not enough validation data for {} GPUs. '
                             'Received {} validation sample.'.format(
                                 test_dict['y'].shape[0], num_gpus))

        # When using multiple GPUs and skip_connections,
        # the training data must be evenly distributed across all GPUs
        num_train = train_dict['y'].shape[0]
        nb_samples = num_train - num_train % batch_size
        if nb_samples:
            train_dict['y'] = train_dict['y'][:nb_samples]
            train_dict['X'] = train_dict['X'][:nb_samples]

    # this will do preprocessing and realtime data augmentation
    datagen = image_generators.RetinaNetGenerator(
        # fill_mode='constant',  # for rotations
        rotation_range=rotation_range,
        shear_range=shear,
        zoom_range=zoom_range,
        horizontal_flip=flip,
        vertical_flip=flip)

    datagen_val = image_generators.RetinaNetGenerator(
        # fill_mode='constant',  # for rotations
        rotation_range=0,
        shear_range=0,
        zoom_range=0,
        horizontal_flip=0,
        vertical_flip=0)

    if 'vgg' in backbone or 'densenet' in backbone:
        compute_shapes = make_shapes_callback(model)
    else:
        compute_shapes = guess_shapes

    train_data = datagen.flow(train_dict,
                              include_masks=include_masks,
                              compute_shapes=compute_shapes,
                              batch_size=batch_size)

    val_data = datagen_val.flow(test_dict,
                                include_masks=include_masks,
                                compute_shapes=compute_shapes,
                                batch_size=batch_size)

    tensorboard_callback = callbacks.TensorBoard(
        log_dir=os.path.join(log_dir, model_name))

    # fit the model on the batches generated by datagen.flow()
    loss_history = model.fit_generator(
        train_data,
        steps_per_epoch=train_data.y.shape[0] // batch_size,
        epochs=n_epoch,
        validation_data=val_data,
        validation_steps=val_data.y.shape[0] // batch_size,
        callbacks=[
            callbacks.LearningRateScheduler(lr_sched),
            callbacks.ModelCheckpoint(model_path,
                                      monitor='val_loss',
                                      verbose=1,
                                      save_best_only=True,
                                      save_weights_only=num_gpus >= 2),
            tensorboard_callback,
            callbacks.ReduceLROnPlateau(monitor='loss',
                                        factor=0.1,
                                        patience=10,
                                        verbose=1,
                                        mode='auto',
                                        min_delta=0.0001,
                                        cooldown=0,
                                        min_lr=0),
            RedirectModel(
                Evaluate(val_data,
                         iou_threshold=iou_threshold,
                         score_threshold=score_threshold,
                         max_detections=max_detections,
                         tensorboard=tensorboard_callback,
                         weighted_average=weighted_average), prediction_model),
        ])

    model.save_weights(model_path)
    np.savez(loss_path, loss_history=loss_history.history)

    average_precisions = evaluate(
        val_data,
        prediction_model,
        iou_threshold=iou_threshold,
        score_threshold=score_threshold,
        max_detections=max_detections,
    )

    # print evaluation
    total_instances = []
    precisions = []
    for label, (average_precision,
                num_annotations) in average_precisions.items():
        print('{:.0f} instances of class'.format(num_annotations), label,
              'with average precision: {:.4f}'.format(average_precision))
        total_instances.append(num_annotations)
        precisions.append(average_precision)

    if sum(total_instances) == 0:
        print('No test instances found.')
    else:
        print(
            'mAP using the weighted average of precisions among classes: {:.4f}'
            .format(
                sum([a * b for a, b in zip(total_instances, precisions)]) /
                sum(total_instances)))
        print('mAP: {:.4f}'.format(
            sum(precisions) / sum(x > 0 for x in total_instances)))

    return model