Ejemplo n.º 1
0
    def _predict_inventory(self, inventory, output_file,
                           batch_size, export_type):
        """ Predict an inventory """
        _ids = list(inventory.keys())
        n_records = len(_ids)
        # process 5 batches at the same time before writing to disk
        n_inventories = calc_n_batches_per_epoch(
            n_records, batch_size * 5, False)
        slices = slice_generator(n_records, n_inventories)
        n_processed = 0
        start_time = time.time()

        # iterate over the dataset
        for i_start, i_end in slices:
            sub_inventory = {k: inventory[k] for k in _ids[i_start:i_end]}

            # create a dataset for the current sub-inventory
            dataset = self._create_dataset_from_inventory(
                sub_inventory, batch_size)
            sub_preds = self._iterate_inventory_dataset(dataset, sub_inventory)
            n_preds = len(sub_preds.keys())

            # export preds
            if export_type == 'csv':
                if i_start == 0:
                    meta_data = set()
                    for k, v in sub_preds.items():
                        if 'meta_data' in v:
                            meta_data = meta_data.union(v['meta_data'].keys())
                    self._create_csv_file_with_header(
                        output_file,
                        meta_headers=list(meta_data))
                self._append_predictions_to_csv(sub_preds, output_file)
            elif export_type == 'json':
                if i_start == 0:
                    self._append_predictions_to_json(
                        output_file, sub_preds,
                        append=False)
                else:
                    self._append_predictions_to_json(
                        output_file, sub_preds,
                        append=True)
            else:
                raise NotImplementedError("export type %s not implemented"
                                          % export_type)

            n_processed += n_preds
            if n_records is not None:
                print_progress(n_processed, n_records)
                estt = estimate_remaining_time(
                    start_time, n_records, n_processed)
                print("\nEstimated time remaining: %s" % estt)
            else:
                print("Processed %s images" % n_processed)

        print("\nProcessed %s records" % n_processed)

        if export_type == 'json':
            self._finish_json_file(output_file)
    def testBoundary(self):
        n_total = 450
        batch_size = 451
        self.assertEqual(
            calc_n_batches_per_epoch(n_total, batch_size,
                                     drop_remainder=False), 1)

        n_total = 450
        batch_size = 449
        self.assertEqual(
            calc_n_batches_per_epoch(n_total, batch_size,
                                     drop_remainder=False), 2)

        n_total = 0
        batch_size = 10
        self.assertEqual(
            calc_n_batches_per_epoch(n_total, batch_size,
                                     drop_remainder=False), 0)
Ejemplo n.º 3
0
def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-train_tfr_path",
        type=str,
        required=True,
        help="Path to directory that contains the training TFR files \
              (incl. subdirs)")
    parser.add_argument(
        "-train_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the training TFR files (default train) \
              list of 1 or more patterns that all have to match",
        default=['train'],
        required=False)
    parser.add_argument(
        "-val_tfr_path",
        type=str,
        required=True,
        help="Path to directory that contains the validation TFR files \
              (incl. subdirs)")
    parser.add_argument(
        "-val_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the validation TFR files (default val) \
              list of 1 or more patterns that all have to match",
        default=['val'],
        required=False)
    parser.add_argument(
        "-test_tfr_path",
        type=str,
        required=False,
        help="Path to directory that contains the test TFR files \
              (incl. subdirs - optional)")
    parser.add_argument(
        "-test_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the test TFR files (default test) \
              list of 1 or more patterns that all have to match",
        default=['test'],
        required=False)
    parser.add_argument(
        "-class_mapping_json",
        type=str,
        required=True,
        help='Path to the json file containing the class mappings')
    parser.add_argument(
        "-run_outputs_dir",
        type=str,
        required=True,
        help="Path to a directory to store data during the training")
    parser.add_argument(
        "-log_outdir",
        type=str,
        required=False,
        default=None,
        help="Directory to write logfiles to (defaults to run_outputs_dir)")
    parser.add_argument("-model_save_dir",
                        type=str,
                        required=True,
                        help='Path to a directory to store final model files')
    parser.add_argument("-model",
                        type=str,
                        required=True,
                        help="The model architecture to use for training\
             (see config/models.yaml)")
    parser.add_argument("-labels",
                        nargs='+',
                        type=str,
                        required=True,
                        help='The labels to model')
    parser.add_argument(
        "-labels_loss_weights",
        nargs='+',
        type=float,
        default=None,
        help='A list of length labels indicating weights for the different\
              labels applied during model training')
    parser.add_argument(
        "-batch_size",
        type=int,
        default=128,
        help="The batch size for model training, if too large the model may\
              crash with an OOM error. Use values between 64 and 256")
    parser.add_argument(
        "-n_cpus",
        type=int,
        default=4,
        help="The number of cpus to use. Use all available if possible.")
    parser.add_argument("-n_gpus",
                        type=int,
                        default=1,
                        help='The number of GPUs to use (default 1)')
    parser.add_argument(
        "-buffer_size",
        type=int,
        default=32768,
        help='The buffer size to use for shuffling training records. Use \
              smaller values if memory is limited.')
    parser.add_argument(
        "-n_parallel_file_reads",
        type=int,
        default=50,
        help='How many files to read in parallel when counting the number of \
              records in tfr files.')
    parser.add_argument("-max_epochs",
                        type=int,
                        default=70,
                        help="The max number of epochs to train the model")
    parser.add_argument("-starting_epoch",
                        type=int,
                        default=0,
                        help="The starting epoch number (0-based index).")
    parser.add_argument(
        "-n_batches_per_epoch_train",
        type=int,
        default=None,
        help="Override automatic calculation of how many batches make up \
              an epoch.")
    ######################################################################
    # Model Training Parameters
    ######################################################################
    parser.add_argument("-initial_learning_rate",
                        type=float,
                        default=0.01,
                        help="The initial learning rate.")
    parser.add_argument("-optimizer",
                        type=str,
                        default="sgd",
                        choices=['sgd', 'rmsprop'],
                        required=False,
                        help="Which optimizer to use in training the model.")
    parser.add_argument(
        "-early_stopping_patience",
        type=int,
        default=3,
        help="Number of epochs after which to stop training if no improvement \
              on validation set was observed (total loss).")
    parser.add_argument(
        "-reduce_lr_on_plateau_patience",
        type=int,
        default=2,
        help="Number of epochs after which to reduce learning rate if no \
              improvement on validation set was observed (total loss).")
    ######################################################################
    # Transfer-Learning and Model Loading
    ######################################################################
    parser.add_argument(
        "-transfer_learning",
        default=False,
        action='store_true',
        required=False,
        help="Option to specify that transfer learning should be used.")
    parser.add_argument(
        "-transfer_learning_type",
        default='last_layer',
        required=False,
        choices=['last_layer', 'all_layers'],
        help="Option to specify which transfer-learning stype hould be used: \
              'last_layer': allowing to adapt only the last layer,  \
              'all_layers': adapt all layers - default is 'last_layer'")
    parser.add_argument(
        "-continue_training",
        default=False,
        action='store_true',
        required=False,
        help="Flag that training should be continued from a saved model.")
    parser.add_argument(
        "-rebuild_model",
        default=False,
        action='store_true',
        required=False,
        help="Flag that model should be rebuild (if continue_training). \
              This might be necessary if model training should be continued\
              with different options (e.g. no GPUs, or different optimizer)")
    parser.add_argument(
        "-model_to_load",
        type=str,
        required=False,
        default=None,
        help='Path to a model (.hdf5) when either continue_training,\
             or transfer_learning are specified, \
             if a directory is specified, \
             the most recent model in that directory is loaded')
    ######################################################################
    # Image Processing
    ######################################################################
    parser.add_argument(
        "-color_augmentation",
        type=str,
        default=None,
        choices=[None, 'little', 'full_fast', 'full_randomized'],
        required=False,
        help="Which (random) color augmentation to perform during model\
              training - choose one of:\
              [None, 'little', 'full_fast', 'full_randomized']. \
              This can slow down the pre-processing speed and starve the \
              GPU of data. Use None or little/full_fast options if input \
              pipeline is slow. Generally full_randomized is recommended \
              and is usually more than fast enough.")
    parser.add_argument(
        "-preserve_aspect_ratio",
        action='store_true',
        default=None,
        dest='preserve_aspect_ratio',
        help="Wheter to preserve the aspect ratio of the images during model \
              training. This keeps the aspect ratio intact which may improve \
              model performance, however, may lead to cut-off areas at the \
              border of an image during training and prediction. If objects \
              of interest may occur at the edges we don't recommend to \
              specify this.")
    parser.add_argument(
        "-ignore_aspect_ratio",
        action='store_false',
        default=None,
        dest='preserve_aspect_ratio',
        help="Wheter to ignore the aspect ratio of the images during model \
              training.")
    parser.add_argument(
        "-randomly_flip_horizontally",
        dest='randomly_flip_horizontally',
        action='store_true',
        default=None,
        help="Whether to randomly flip the image during model training. \
              This almost always makes sense unless the training labels \
              are not invariant to flipping.")
    parser.add_argument(
        "-dont_randomly_flip_horizontally",
        dest='randomly_flip_horizontally',
        action='store_false',
        default=None,
        help="Whether to not randomly flip the image during model training. \
              This only makes sense if the training labels \
              are not invariant to flipping.")
    parser.add_argument("-crop_factor",
                        type=float,
                        default=None,
                        metavar="[0-0.5]",
                        help="Wheter to randomly crop the image within \
             image-size * [1-crop_factor, 1] during model training. \
             This extracts a random portion of the image. \
             Values between 0 and 0.5 are allowed.")
    parser.add_argument("-zoom_factor",
                        type=float,
                        default=None,
                        metavar="[0-0.5]",
                        help="Wheter to randomly zoom the image by: \
             image-size * [1 -zoom_factor, 1+zoom_factor] during model \
             training. Values between 0 and 0.5 are allowed.")
    parser.add_argument(
        "-rotate_by_angle",
        type=int,
        default=None,
        choices=range(0, 181),
        metavar="[0-180]",
        help="Wheter to randomly rotate the image during model training in \
              range [0 - rotate_by_angle, 0 + rotate_by_angle] \
             training. Values between 0 and 180 degrees are allowed.")
    parser.add_argument(
        "-image_choice_for_sets",
        type=str,
        default=None,
        choices=['random', 'grayscale_stacking'],
        help="How to choose an image for records with multiple images. \
              Default is 'random' which randomly chooses an image \
              during model training. 'grayscale_stacking' converts multiple \
              images into a single RGB image by blurring and converting \
              individual images to grayscale. Note that grayscale_stacking is \
              an experimental feature and is not yet supported when using \
              the predictor on new images. ")
    parser.add_argument(
        "-output_width",
        type=int,
        default=None,
        help="The output width in pixels of the image after pre-processing, \
              thus the input width of the image into the model. \
              Use this to override the default model-specific values as \
              specified in the config.yaml")
    parser.add_argument(
        "-output_height",
        type=int,
        default=None,
        help="The output height in pixels of the image after pre-processing, \
              thus the input height of the image into the model. \
              Use this to override the default model-specific values as \
              specified in the config.yaml")

    # Parse command line arguments
    args = vars(parser.parse_args())

    # Configure Logging
    if args['log_outdir'] is None:
        args['log_outdir'] = args['run_outputs_dir']

    setup_logging(log_output_path=args['log_outdir'])

    logger = logging.getLogger(__name__)

    logger.info("Using arguments:")
    for k, v in args.items():
        logger.info("Arg: %s: %s" % (k, v))

    ###########################################
    # Process Input ###########
    ###########################################

    # Load config file
    cfg_path = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                            'config', 'config.yaml')

    config = ConfigLoader(cfg_path)

    assert args['model'] in config.cfg['models'], \
        "model %s not found in config/models.yaml" % args['model']

    # get default image_processing options from config
    image_processing = config.cfg['image_processing']

    # load model specific image processing parameters
    image_processing_model = \
        config.cfg['models'][args['model']]['image_processing']

    # overwrite parameters if specified by user
    to_overwrite = [
        'color_augmentation', 'preserve_aspect_ratio', 'crop_factor',
        'zoom_factor', 'rotate_by_angle', 'randomly_flip_horizontally',
        'image_choice_for_sets', 'output_width', 'output_height'
    ]
    for overwrite in to_overwrite:
        if args[overwrite] is not None:
            image_processing[overwrite] = args[overwrite]

    image_processing = {**image_processing_model, **image_processing}

    # disable color_augmentation for grayscale_stacking
    if image_processing['image_choice_for_sets'] == 'grayscale_stacking':
        if image_processing['color_augmentation'] is not None:
            image_processing['color_augmentation'] = None
            msg = "Disabling color_augmentation because of \
                   incompatibility with grayscale_stacking"

            logger.info(textwrap.shorten(msg, width=99))

    input_shape = (image_processing['output_height'],
                   image_processing['output_width'], 3)

    # Add 'label/' prefix to labels as they are stored in the .tfrecord files
    output_labels = args['labels']
    output_labels_clean = ['label/' + x for x in output_labels]

    # Class to numeric mappings and number of classes per label
    class_mapping = read_json(args['class_mapping_json'])
    # TODO: fix num classes per label for a:0, b:0 cases
    n_classes_per_label_dict = {
        c: len(set(class_mapping[o].values()))
        for o, c in zip(output_labels, output_labels_clean)
    }
    n_classes_per_label = [
        n_classes_per_label_dict[x] for x in output_labels_clean
    ]

    # save class mapping file to current run path
    export_dict_to_json(
        class_mapping,
        os.path.join(args['run_outputs_dir'], 'label_mappings.json'))

    # Find TFR files
    tfr_train = find_tfr_files_pattern_subdir(args['train_tfr_path'],
                                              args['train_tfr_pattern'])
    tfr_val = find_tfr_files_pattern_subdir(args['val_tfr_path'],
                                            args['val_tfr_pattern'])

    # Create best model output name
    best_model_save_path = os.path.join(args['model_save_dir'],
                                        'best_model.hdf5')

    # Define path of model to load if only directory is specified
    if args['model_to_load'] is not None:
        if not args['model_to_load'].endswith('.hdf5'):
            if os.path.isdir(args['model_to_load']):
                model_files = \
                    find_files_with_ending(args['model_to_load'], '.hdf5')
                most_recent_model = \
                    get_most_recent_file_from_files(model_files)
                args['model_to_load'] = most_recent_model
                logger.debug("Loading most recent model file %s:" %
                             most_recent_model)

    ###########################################
    # CALC IMAGE STATS ###########
    ###########################################

    logger.info("Start Calculating Image Stats")

    tfr_encoder_decoder = DefaultTFRecordEncoderDecoder()
    data_reader = DatasetReader(tfr_encoder_decoder.decode_record)

    # Calculate Dataset Image Means and Stdevs for a dummy batch
    logger.info("Get Dataset Reader for calculating dataset stats")
    n_records_train = n_records_in_tfr_dataset(
        tfr_train, n_parallel_file_reads=args['n_parallel_file_reads'])
    dataset = data_reader.get_iterator(
        tfr_files=tfr_train,
        batch_size=min([4096, n_records_train]),
        is_train=True,
        n_repeats=1,
        output_labels=output_labels,
        image_pre_processing_fun=preprocess_image,
        image_pre_processing_args={
            **image_processing, 'is_training': False
        },
        buffer_size=min([args['buffer_size'], 128]),
        num_parallel_calls=args['n_cpus'])
    iterator = dataset.make_one_shot_iterator()
    batch_data = iterator.get_next()

    logger.info("Calculating image means and stdevs")
    with tf.Session() as sess:
        features, labels = sess.run(batch_data)

    # calculate and save image means and stdvs of each color channel
    # for pre processing purposes
    image_means = [
        round(float(x), 4) for x in list(
            np.mean(features['images'], axis=(0, 1, 2), dtype=np.float64))
    ]
    image_stdevs = [
        round(float(x), 4) for x in list(
            np.std(features['images'], axis=(0, 1, 2), dtype=np.float64))
    ]

    image_processing['image_means'] = image_means
    image_processing['image_stdevs'] = image_stdevs

    logger.info("Image Means: %s" % image_means)
    logger.info("Image Stdevs: %s" % image_stdevs)

    # Export Image Processing Settings
    export_dict_to_json({
        **image_processing, 'is_training': False
    }, os.path.join(args['run_outputs_dir'], 'image_processing.json'))

    ###########################################
    # PREPARE DATA READER ###########
    ###########################################

    logger.info("Preparing Data Feeders")

    def input_feeder_train():
        return data_reader.get_iterator(
            tfr_files=tfr_train,
            batch_size=args['batch_size'],
            is_train=True,
            n_repeats=None,
            output_labels=output_labels,
            label_to_numeric_mapping=class_mapping,
            image_pre_processing_fun=preprocess_image,
            image_pre_processing_args={
                **image_processing, 'is_training': True
            },
            buffer_size=args['buffer_size'],
            num_parallel_calls=args['n_cpus'])

    def input_feeder_val():
        return data_reader.get_iterator(
            tfr_files=tfr_val,
            batch_size=args['batch_size'],
            is_train=False,
            n_repeats=None,
            output_labels=output_labels,
            label_to_numeric_mapping=class_mapping,
            image_pre_processing_fun=preprocess_image,
            image_pre_processing_args={
                **image_processing, 'is_training': False
            },
            buffer_size=args['buffer_size'],
            num_parallel_calls=args['n_cpus'])

    logger.info("Calculating batches per epoch")
    if args['n_batches_per_epoch_train'] is None:
        n_batches_per_epoch_train = calc_n_batches_per_epoch(
            n_records_train, args['batch_size'])
    else:
        n_batches_per_epoch_train = args['n_batches_per_epoch_train']

    n_records_val = n_records_in_tfr_dataset(
        tfr_val, n_parallel_file_reads=args['n_parallel_file_reads'])
    n_batches_per_epoch_val = calc_n_batches_per_epoch(n_records_val,
                                                       args['batch_size'])

    logger.info("Found %s records in the training set" % n_records_train)
    logger.debug("Using %s batches/epoch for the training set" %
                 n_batches_per_epoch_train)
    logger.info("Found %s records in the validation set" % n_records_val)
    logger.debug("Using %s batches/epoch for the validation set" %
                 n_batches_per_epoch_val)

    ###########################################
    # CREATE MODEL ###########
    ###########################################

    logger.info("Preparing Model")

    model = create_model(model_name=args['model'],
                         input_shape=input_shape,
                         target_labels=output_labels_clean,
                         n_classes_per_label_type=n_classes_per_label,
                         n_gpus=args['n_gpus'],
                         continue_training=args['continue_training'],
                         rebuild_model=args['rebuild_model'],
                         transfer_learning=args['transfer_learning'],
                         transfer_learning_type=args['transfer_learning_type'],
                         path_of_model_to_load=args['model_to_load'],
                         initial_learning_rate=args['initial_learning_rate'],
                         output_loss_weights=args['labels_loss_weights'])

    logger.debug("Final Model Architecture")
    for layer, i in zip(model.layers, range(0, len(model.layers))):
        logger.debug("Layer %s: Name: %s Input: %s Output: %s" %
                     (i, layer.name, layer.input_shape, layer.output_shape))

    ###########################################
    # MONITORS / HOOKS ###########
    ###########################################

    logger.info("Preparing Callbacks and Monitors")

    # stop model training if validation loss does not improve
    early_stopping = EarlyStopping(monitor='val_loss',
                                   min_delta=0,
                                   patience=args['early_stopping_patience'],
                                   verbose=0,
                                   mode='auto')

    # reduce learning rate if model progress plateaus
    reduce_lr_on_plateau = ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.1,
        patience=args['reduce_lr_on_plateau_patience'],
        verbose=0,
        mode='auto',
        min_delta=0.0001,
        cooldown=1,
        min_lr=1e-5)

    # log validation statistics to a csv file
    csv_logger = CSVLogger(args['run_outputs_dir'] + 'training.log',
                           append=args['continue_training'])

    # Log metrics after each epoch to logger
    logging_logger = LoggingLogger(logger)

    # create model checkpoints after each epoch
    checkpointer = ModelCheckpoint(
        filepath=args['run_outputs_dir'] +
        'model_epoch_{epoch:02d}_loss_{val_loss:.2f}.hdf5',
        monitor='val_loss',
        verbose=0,
        save_best_only=False,
        save_weights_only=False,
        mode='auto',
        period=1)

    # save best model
    checkpointer_best = ModelCheckpoint(filepath=args['run_outputs_dir'] +
                                        'model_best.hdf5',
                                        monitor='val_loss',
                                        verbose=0,
                                        save_best_only=True,
                                        save_weights_only=False,
                                        mode='auto',
                                        period=1)

    # write graph to disk
    tensorboard = TensorBoard(log_dir=args['run_outputs_dir'],
                              histogram_freq=0,
                              batch_size=args['batch_size'],
                              write_graph=True,
                              write_grads=False,
                              write_images=False)

    # Initialize tables (lookup tables)
    table_init = TableInitializerCallback()

    callbacks_list = [
        early_stopping, reduce_lr_on_plateau, csv_logger, logging_logger,
        checkpointer, checkpointer_best, table_init, tensorboard
    ]

    ###########################################
    # MODEL TRAINING  ###########
    ###########################################

    logger.info("Start Model Training")

    model.fit(input_feeder_train(),
              epochs=args['max_epochs'],
              steps_per_epoch=n_batches_per_epoch_train,
              validation_data=input_feeder_val(),
              validation_steps=n_batches_per_epoch_val,
              callbacks=callbacks_list,
              initial_epoch=args['starting_epoch'])

    logger.info("Finished Model Training")

    ###########################################
    # SAVE BEST MODEL ###########
    ###########################################

    logger.info("Saving Best Model")

    copy_models_and_config_files(model_source=args['run_outputs_dir'] +
                                 'model_best.hdf5',
                                 model_target=best_model_save_path,
                                 files_path_source=args['run_outputs_dir'],
                                 files_path_target=args['model_save_dir'],
                                 copy_files=".json")

    ###########################################
    # PREDICT AND EXPORT TEST DATA ###########
    ###########################################

    if len(args['test_tfr_path']) > 0:

        logger.info("Starting to predict on test data")

        tfr_test = find_tfr_files_pattern_subdir(args['test_tfr_path'],
                                                 args['test_tfr_pattern'])

        pred_output_json = os.path.join(args['run_outputs_dir'],
                                        'test_preds.json')

        tf.keras.backend.clear_session()

        tfr_encoder_decoder = DefaultTFRecordEncoderDecoder()
        logger.info("Create Dataset Reader")
        data_reader = DatasetReader(tfr_encoder_decoder.decode_record)

        def input_feeder_test():
            return data_reader.get_iterator(
                tfr_files=tfr_test,
                batch_size=args['batch_size'],
                is_train=False,
                n_repeats=1,
                output_labels=output_labels,
                image_pre_processing_fun=preprocess_image,
                image_pre_processing_args={
                    **image_processing, 'is_training': False
                },
                buffer_size=args['buffer_size'],
                num_parallel_calls=args['n_cpus'],
                drop_batch_remainder=False)

        pred = Predictor(model_path=best_model_save_path,
                         class_mapping_json=args['class_mapping_json'],
                         pre_processing_json=args['run_outputs_dir'] +
                         'image_processing.json')

        pred.predict_from_dataset(dataset=input_feeder_test(),
                                  export_type='json',
                                  output_file=pred_output_json)

        logger.info("Finished predicting on test data, saved to: %s" %
                    pred_output_json)
Ejemplo n.º 4
0
def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-train_tfr_path",
        type=str,
        required=True,
        help="Path to directory that contains the training TFR files \
              (incl. subdirs)")
    parser.add_argument(
        "-train_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the training TFR files (default train) \
              list of 1 or more patterns that all have to match",
        default=['train'],
        required=False)
    parser.add_argument(
        "-val_tfr_path",
        type=str,
        required=True,
        help="Path to directory that contains the validation TFR files \
              (incl. subdirs)")
    parser.add_argument(
        "-val_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the validation TFR files (default val) \
              list of 1 or more patterns that all have to match",
        default=['val'],
        required=False)
    parser.add_argument(
        "-test_tfr_path",
        type=str,
        required=False,
        help="Path to directory that contains the test TFR files \
              (incl. subdirs - optional)")
    parser.add_argument(
        "-test_tfr_pattern",
        nargs='+',
        type=str,
        help="The pattern of the test TFR files (default test) \
              list of 1 or more patterns that all have to match",
        default=['test'],
        required=False)
    parser.add_argument(
        "-class_mapping_json",
        type=str,
        required=True,
        help='Path to the json file containing the class mappings')
    parser.add_argument(
        "-run_outputs_dir",
        type=str,
        required=True,
        help="Path to a directory to store data during the training")
    parser.add_argument(
        "-log_outdir",
        type=str,
        required=False,
        default=None,
        help="Directory to write logfiles to (defaults to run_outputs_dir)")
    parser.add_argument("-model_save_dir",
                        type=str,
                        required=True,
                        help='Path to a directory to store final model files')
    parser.add_argument("-model",
                        type=str,
                        required=True,
                        help="The model architecture to use for training\
             (see config/models.yaml)")
    parser.add_argument("-labels",
                        nargs='+',
                        type=str,
                        required=True,
                        help='The labels to model')
    parser.add_argument(
        "-labels_loss_weights",
        nargs='+',
        type=float,
        default=None,
        help='A list of length labels indicating weights for the different\
              labels applied during model training')
    parser.add_argument(
        "-batch_size",
        type=int,
        default=128,
        help="The batch size for model training, if too large the model may\
              crash with an OOM error. Use values between 64 and 256")
    parser.add_argument(
        "-n_cpus",
        type=int,
        default=4,
        help="The number of cpus to use. Use all available if possible.")
    parser.add_argument("-n_gpus",
                        type=int,
                        default=1,
                        help='The number of GPUs to use (default 1)')
    parser.add_argument(
        "-buffer_size",
        type=int,
        default=32768,
        help='The buffer size to use for shuffling training records. Use \
              smaller values if memory is limited.')
    parser.add_argument("-max_epochs",
                        type=int,
                        default=70,
                        help="The max number of epochs to train the model")
    parser.add_argument("-starting_epoch",
                        type=int,
                        default=0,
                        help="The starting epoch number (0-based index).")
    # Model Training Parameters
    parser.add_argument("-initial_learning_rate",
                        type=float,
                        default=0.01,
                        help="The initial learning rate.")
    parser.add_argument(
        "-optimizer",
        type=str,
        default="sgd",
        required=False,
        help="Which optimizer to use in training the model (sgd or rmsprop)")
    # Transfer-Learning and Model Loading
    parser.add_argument(
        "-transfer_learning",
        default=False,
        action='store_true',
        required=False,
        help="Option to specify that transfer learning should be used.")
    parser.add_argument(
        "-transfer_learning_type",
        default='last_layer',
        required=False,
        help="Option to specify that transfer learning should be used, by\
              allowing to adapt only the last layer ('last_layer') \
              or all layers ('all_layers') - default is 'last_layer'")
    parser.add_argument(
        "-continue_training",
        default=False,
        action='store_true',
        required=False,
        help="Flag that training should be continued from a saved model.")
    parser.add_argument(
        "-rebuild_model",
        default=False,
        action='store_true',
        required=False,
        help="Flag that model should be rebuild (if continue_training). \
              This might be necessary if model training should be continued\
              with different options (e.g. no GPUs, or different optimizer)")
    parser.add_argument(
        "-model_to_load",
        type=str,
        required=False,
        default=None,
        help='Path to a model (.hdf5) when either continue_training,\
             or transfer_learning are specified, \
             if a directory is specified, \
             the most recent model in that directory is loaded')
    # Image Processing
    parser.add_argument(
        "-color_augmentation",
        type=str,
        default="full_randomized",
        required=False,
        help="Which (random) color augmentation to perform during model\
              training - choose one of:\
              [None, 'little', 'full_fast', 'full_randomized']. \
              This can slow down the pre-processing speed and starve the \
              GPU of data. Use None or little/full_fast options if input \
              pipeline is slow. Generally full_randomized is recommended \
              and is usually more than fast enough.")
    parser.add_argument(
        "-ignore_aspect_ratio",
        default=False,
        action='store_true',
        help="Wheter to ignore the aspect ratio of the images during model \
              training. This can improve the total area of the image the \
              model sees during training and prediction. However, the images \
              are slightly distorted with this option since they are \
              converted to squares.")

    # Parse command line arguments
    args = vars(parser.parse_args())

    # Configure Logging
    if args['log_outdir'] is None:
        args['log_outdir'] = args['run_outputs_dir']

    setup_logging(log_output_path=args['log_outdir'])

    logger = logging.getLogger(__name__)

    print("Using arguments:")
    for k, v in args.items():
        print("Arg: %s: %s" % (k, v))

    ###########################################
    # Process Input ###########
    ###########################################

    # Load model config
    models_cfg_path = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                                   'config', 'models.yaml')

    model_cfg = ConfigLoader(models_cfg_path)

    assert args['model'] in model_cfg.cfg['models'], \
        "model %s not found in config/models.yaml" % args['model']

    image_processing = \
        model_cfg.cfg['models'][args['model']]['image_processing']
    image_processing['ignore_aspect_ratio'] = args['ignore_aspect_ratio']

    input_shape = (image_processing['output_height'],
                   image_processing['output_width'], 3)

    # Add 'label/' prefix to labels as they are stored in the .tfrecord files
    output_labels = args['labels']
    output_labels_clean = ['label/' + x for x in output_labels]

    # Class to numeric mappings and number of classes per label
    class_mapping = read_json(args['class_mapping_json'])
    # TODO: fix num classes per label for a:0, b:0 cases
    n_classes_per_label_dict = {
        c: len(set(class_mapping[o].values()))
        for o, c in zip(output_labels, output_labels_clean)
    }
    n_classes_per_label = [
        n_classes_per_label_dict[x] for x in output_labels_clean
    ]

    # save class mapping file to current run path
    export_dict_to_json(
        class_mapping,
        os.path.join(args['run_outputs_dir'], 'label_mappings.json'))

    # Find TFR files
    tfr_train = find_tfr_files_pattern_subdir(args['train_tfr_path'],
                                              args['train_tfr_pattern'])
    tfr_val = find_tfr_files_pattern_subdir(args['val_tfr_path'],
                                            args['val_tfr_pattern'])

    if len(args['test_tfr_path']) > 0:
        TEST_SET = True
        tfr_test = find_tfr_files_pattern_subdir(args['test_tfr_path'],
                                                 args['test_tfr_pattern'])
        pred_output_json = os.path.join(args['run_outputs_dir'],
                                        'test_preds.json')
    else:
        TEST_SET = False

    # Create best model output name
    best_model_save_path = os.path.join(args['model_save_dir'],
                                        'best_model.hdf5')

    # Define path of model to load if only directory is specified
    if args['model_to_load'] is not None:
        if not args['model_to_load'].endswith('.hdf5'):
            if os.path.isdir(args['model_to_load']):
                model_files = \
                    find_files_with_ending(args['model_to_load'], '.hdf5')
                most_recent_model = \
                    get_most_recent_file_from_files(model_files)
                args['model_to_load'] = most_recent_model
                logging.debug("Loading most recent model file %s:" %
                              most_recent_model)

    ###########################################
    # CALC IMAGE STATS ###########
    ###########################################

    tfr_encoder_decoder = DefaultTFRecordEncoderDecoder()

    logger.info("Create Dataset Reader")
    data_reader = DatasetReader(tfr_encoder_decoder.decode_record)

    # Calculate Dataset Image Means and Stdevs for a dummy batch
    logger.info("Get Dataset Reader for calculating datset stats")
    n_records_train = n_records_in_tfr_parallel(tfr_train, args['n_cpus'])
    dataset = data_reader.get_iterator(
        tfr_files=tfr_train,
        batch_size=min([4096, n_records_train]),
        is_train=True,
        n_repeats=1,
        output_labels=output_labels,
        label_to_numeric_mapping=class_mapping,
        image_pre_processing_fun=preprocess_image,
        image_pre_processing_args={
            **image_processing, 'is_training': False
        },
        buffer_size=args['buffer_size'],
        num_parallel_calls=args['n_cpus'])
    iterator = dataset.make_initializable_iterator()
    batch_data = iterator.get_next()

    logger.info("Calculating image means and stdevs")
    with tf.Session() as sess:
        tf.tables_initializer().run()
        sess.run(iterator.initializer)
        features, labels = sess.run(batch_data)

    # calculate and save image means and stdvs of each color channel
    # for pre processing purposes
    image_means = [
        round(float(x), 4) for x in list(
            np.mean(features['images'], axis=(0, 1, 2), dtype=np.float64))
    ]
    image_stdevs = [
        round(float(x), 4) for x in list(
            np.std(features['images'], axis=(0, 1, 2), dtype=np.float64))
    ]

    image_processing['image_means'] = image_means
    image_processing['image_stdevs'] = image_stdevs

    logger.info("Image Means: %s" % image_means)
    logger.info("Image Stdevs: %s" % image_stdevs)

    ###########################################
    # PREPARE DATA READER ###########
    ###########################################

    logger.info("Preparing Data Feeders")

    def input_feeder_train():
        return data_reader.get_iterator(
            tfr_files=tfr_train,
            batch_size=args['batch_size'],
            is_train=True,
            n_repeats=None,
            output_labels=output_labels,
            label_to_numeric_mapping=class_mapping,
            image_pre_processing_fun=preprocess_image,
            image_pre_processing_args={
                **image_processing, 'is_training': True,
                'color_augmentation': args['color_augmentation']
            },
            buffer_size=args['buffer_size'],
            num_parallel_calls=args['n_cpus'])

    def input_feeder_val():
        return data_reader.get_iterator(
            tfr_files=tfr_val,
            batch_size=args['batch_size'],
            is_train=False,
            n_repeats=None,
            output_labels=output_labels,
            label_to_numeric_mapping=class_mapping,
            image_pre_processing_fun=preprocess_image,
            image_pre_processing_args={
                **image_processing, 'is_training': False
            },
            buffer_size=args['buffer_size'],
            num_parallel_calls=args['n_cpus'])

    if TEST_SET:

        def input_feeder_test():
            return data_reader.get_iterator(
                tfr_files=tfr_test,
                batch_size=args['batch_size'],
                is_train=False,
                n_repeats=1,
                output_labels=output_labels,
                image_pre_processing_fun=preprocess_image,
                image_pre_processing_args={
                    **image_processing, 'is_training': False
                },
                buffer_size=args['buffer_size'],
                num_parallel_calls=args['n_cpus'],
                drop_batch_remainder=False)

    # Export Image Processing Settings
    export_dict_to_json({
        **image_processing, 'is_training': False
    }, os.path.join(args['run_outputs_dir'], 'image_processing.json'))

    logger.info("Calculating batches per epoch")
    n_batches_per_epoch_train = calc_n_batches_per_epoch(
        n_records_train, args['batch_size'])

    n_records_val = n_records_in_tfr_parallel(tfr_val, args['n_cpus'])
    n_batches_per_epoch_val = calc_n_batches_per_epoch(n_records_val,
                                                       args['batch_size'])

    ###########################################
    # CREATE MODEL ###########
    ###########################################

    logger.info("Preparing Model")

    model = create_model(model_name=args['model'],
                         input_shape=input_shape,
                         target_labels=output_labels_clean,
                         n_classes_per_label_type=n_classes_per_label,
                         n_gpus=args['n_gpus'],
                         continue_training=args['continue_training'],
                         rebuild_model=args['rebuild_model'],
                         transfer_learning=args['transfer_learning'],
                         transfer_learning_type=args['transfer_learning_type'],
                         path_of_model_to_load=args['model_to_load'],
                         initial_learning_rate=args['initial_learning_rate'],
                         output_loss_weights=args['labels_loss_weights'])

    logger.debug("Final Model Architecture")
    for layer, i in zip(model.layers, range(0, len(model.layers))):
        logger.debug("Layer %s: Name: %s Input: %s Output: %s" %
                     (i, layer.name, layer.input_shape, layer.output_shape))

    logger.info("Preparing Callbacks and Monitors")

    ###########################################
    # MONITORS / HOOKS ###########
    ###########################################

    # stop model training if validation loss does not improve
    early_stopping = EarlyStopping(monitor='val_loss',
                                   min_delta=0,
                                   patience=3,
                                   verbose=0,
                                   mode='auto')

    # reduce learning rate if model progress plateaus
    reduce_lr_on_plateau = ReduceLROnPlateau(monitor='val_loss',
                                             factor=0.1,
                                             patience=2,
                                             verbose=0,
                                             mode='auto',
                                             min_delta=0.0001,
                                             cooldown=1,
                                             min_lr=1e-5)

    # log validation statistics to a csv file
    csv_logger = CSVLogger(args['run_outputs_dir'] + 'training.log',
                           append=args['continue_training'])

    # create model checkpoints after each epoch
    checkpointer = ModelCheckpoint(
        filepath=args['run_outputs_dir'] +
        'model_epoch_{epoch:02d}_loss_{val_loss:.2f}.hdf5',
        monitor='val_loss',
        verbose=0,
        save_best_only=False,
        save_weights_only=False,
        mode='auto',
        period=1)

    # save best model
    checkpointer_best = ModelCheckpoint(filepath=args['run_outputs_dir'] +
                                        'model_best.hdf5',
                                        monitor='val_loss',
                                        verbose=0,
                                        save_best_only=True,
                                        save_weights_only=False,
                                        mode='auto',
                                        period=1)

    # write graph to disk
    tensorboard = TensorBoard(log_dir=args['run_outputs_dir'],
                              histogram_freq=0,
                              batch_size=args['batch_size'],
                              write_graph=True,
                              write_grads=False,
                              write_images=False)

    # Initialize tables (lookup tables)
    table_init = TableInitializerCallback()

    callbacks_list = [
        early_stopping, reduce_lr_on_plateau, csv_logger, checkpointer,
        checkpointer_best, table_init, tensorboard
    ]

    ###########################################
    # MODEL TRAINING  ###########
    ###########################################

    logger.info("Start Model Training")

    model.fit(input_feeder_train(),
              epochs=args['max_epochs'],
              steps_per_epoch=n_batches_per_epoch_train,
              validation_data=input_feeder_val(),
              validation_steps=n_batches_per_epoch_val,
              callbacks=callbacks_list,
              initial_epoch=args['starting_epoch'])

    logger.info("Finished Model Training")

    ###########################################
    # SAVE BEST MODEL ###########
    ###########################################

    copy_models_and_config_files(model_source=args['run_outputs_dir'] +
                                 'model_best.hdf5',
                                 model_target=best_model_save_path,
                                 files_path_source=args['run_outputs_dir'],
                                 files_path_target=args['model_save_dir'],
                                 copy_files=".json")

    ###########################################
    # PREDICT AND EXPORT TEST DATA ###########
    ###########################################

    if TEST_SET:
        logger.info("Starting to predict on test data")
        pred = Predictor(model_path=best_model_save_path,
                         class_mapping_json=args['class_mapping_json'],
                         pre_processing_json=args['run_outputs_dir'] +
                         'image_processing.json',
                         batch_size=args['batch_size'])

        pred.predict_from_dataset(dataset=input_feeder_test(),
                                  export_type='json',
                                  output_file=pred_output_json)

        logger.info("Finished predicting on test data, saved to: %s" %
                    pred_output_json)
 def testEqual(self):
     n_total = 6600
     batch_size = 6600
     self.assertEqual(
         calc_n_batches_per_epoch(n_total, batch_size,
                                  drop_remainder=False), 1)
 def testDropRemainderCaseNormal(self):
     n_total = 26
     batch_size = 5
     self.assertEqual(calc_n_batches_per_epoch(n_total, batch_size), 5)
 def testNormalCase(self):
     n_total = 26
     batch_size = 5
     self.assertEqual(
         calc_n_batches_per_epoch(n_total, batch_size,
                                  drop_remainder=False), 6)