def main(): FILE_DIRECTORY = os.path.dirname(os.path.abspath(__file__)) # Parse arguments parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default=os.path.join(FILE_DIRECTORY, "configs/local.cfg")) args = parser.parse_args() if not os.path.exists(args.config): raise FileNotFoundError("ERROR: Config file \"%s\" not found" % (args.config)) else: cfg = get_config(args.config) model_name = args.config.split('/')[-1] model_name = model_name.split('.')[0] cfg.model_name = cfg.runtime + "_" + model_name assert cfg.inference is False, "You are running a training in inference mode. Please check your config!" # setting seed from .utils.helpers import initialize_seed # Set seed to compare trainings and exclude randomness initialize_seed(cfg.seed) # create object that stores backbone information backbone = architectures.backbone(cfg.network) # make sure keras is the minimum required version check_keras_version() # optionally choose specific GPU os.environ['CUDA_VISIBLE_DEVICES'] = cfg.gpu keras.backend.tensorflow_backend.set_session(get_session( cfg.gpu_mem_usage)) # create the generators if 'nuscenes' in cfg.data_set: train_generator, validation_generator, test_generator, test_night_generator, test_rain_generator = create_generators( cfg, backbone) else: train_generator, validation_generator = create_generators( cfg, backbone) # create the model weights = None if cfg.load_model: print('Loading model, this may take a second...') model = architectures.load_model(cfg.load_model, backbone_name=cfg.network) training_model = model prediction_model = retinanet_bbox( model=model, anchor_params=None, class_specific_filter=cfg.class_specific_nms) else: if cfg.pretrain_basenet: weights = backbone.download_imagenet() in_shape = (cfg.image_size[0], cfg.image_size[1], len(train_generator.channels)) print('Creating model, this may take a second...') model, training_model, prediction_model = create_models( backbone_retinanet=backbone.retinanet, num_classes=train_generator.num_classes(), weights=weights, multi_gpu=0, freeze_backbone=False, lr=cfg.learning_rate, inputs=in_shape, cfg=cfg, distance=cfg.distance_detection, distance_alpha=cfg.distance_alpha) # print model summary print(model.summary()) print("Model Parameters: ", model.count_params()) # this lets the generator compute backbone layer shapes using the actual backbone model if 'vgg' in cfg.network or 'densenet' in cfg.network: train_generator.compute_shapes = make_shapes_callback(model) if validation_generator: validation_generator.compute_shapes = train_generator.compute_shapes # create the callbacks callbacks = create_callbacks( model, prediction_model, validation_generator, cfg, ) # Use multiprocessing if cpu_count > 0 use_multiprocessing = cfg.workers > 0 # class weights class_weights_labels = {} if cfg.class_weights: class_weights_names = cfg.class_weights for key in class_weights_names.keys(): class_weights_labels[train_generator.name_to_label(key)] = float( class_weights_names[key]) # Print outputs print() print("=" * 60) print("\t\t##### Parameters #####") print("=" * 60) descr = cfg.get_description() descr = os.linesep.join([s for s in descr.splitlines() if s.strip()]) print(descr) print() print("=" * 60) print("\t\t##### Start Training #####") print("=" * 60) ## Start training training_model.fit_generator(generator=train_generator, steps_per_epoch=len(train_generator), epochs=cfg.epochs, validation_data=validation_generator, validation_steps=len(validation_generator), verbose=1, callbacks=callbacks, workers=cfg.workers, use_multiprocessing=use_multiprocessing, class_weight=class_weights_labels) ## Evaluate on test data_set print("=" * 60) print("\t\t##### Evaluate Test Set #####") print("=" * 60) # Load best model best_model = keras.models.load_model( cfg.save_model + cfg.model_name + '.h5', custom_objects=backbone.custom_objects) # load anchor parameters, or pass None (so that defaults will be used) if 'small' in cfg.anchor_params: anchor_params = AnchorParameters.small else: anchor_params = None best_prediction_model = retinanet_bbox(model=best_model, anchor_params=anchor_params, class_specific_filter=False) # Evaluate from .utils.eval_test import evaluate_test_set evaluate_test_set(best_prediction_model, test_generator, cfg, mode='all', tensorboard=callbacks[1], verbose=1) print("=" * 60) print("\t##### Evaluate Test Set at Night #####") print("=" * 60) evaluate_test_set(best_prediction_model, test_night_generator, cfg, mode='night', tensorboard=callbacks[1], verbose=1) print("=" * 60) print("\t##### Evaluate Test Set at Rain #####") print("=" * 60) evaluate_test_set(best_prediction_model, test_rain_generator, cfg, mode='rain', tensorboard=callbacks[1], verbose=1) print("=" * 60) print("\t######## Finished successfully ########") print("=" * 60)
data = pickle.load(fp) [all_detections, all_annotations] = data else: # load model model = keras.models.load_model(args.model, custom_objects=backbone.custom_objects) # load anchor parameters, or pass None (so that defaults will be used) if 'small' in cfg.anchor_params: anchor_params = AnchorParameters.small num_anchors = AnchorParameters.small.num_anchors() else: anchor_params = None num_anchors = None prediction_model = retinanet_bbox(model=model, anchor_params=anchor_params, class_specific_filter=False) all_detections = _get_detections(test_generator, prediction_model, distance=cfg.distance_detection, score_threshold=score_threshold, max_detections=100, save_path=None, render=args.render, distance_scale=100, workers=cfg.workers, cfg=cfg) all_annotations = _get_annotations(test_generator) pickle_path = './saved_models/detection_pickles'
def create_models(backbone_retinanet, num_classes, weights, multi_gpu=0, freeze_backbone=False, distance=False, distance_alpha=1.0, lr=1e-5, cfg=None, inputs=(None, None, 3)): """ Creates three models (model, training_model, prediction_model). :param backbone_retinanet: <func> A function to call to create a retinanet model with a given backbone :param num_classes: <int> The number of classes to train :param weights: <keras.Weights> The weights to load into the model :param multi_gpu: <int> The number of GPUs to use for training :param freeze_backbone: <bool> If True, disables learning for the backbone :param distance: <bool> If True, distance detection is enabled :param distance_alpha: <float> Weighted loss factor for distance loss :param lr: <float> Learning rate for network training :param cfg: <Configuration> Config class with config parameters :param inputs: <tuple> Input shape for neural network :return model: <keras.Model> The base model. This is also the model that is saved in snapshots. :return training_model: <keras.Model> The training model. If multi_gpu=0, this is identical to model. :return prediction_model: <keras.Model> The model wrapped with utility functions to perform object detection (applies regression values and performs NMS). """ modifier = freeze_model if freeze_backbone else None # load anchor parameters, or pass None (so that defaults will be used) if 'small' in cfg.anchor_params: anchor_params = AnchorParameters.small num_anchors = AnchorParameters.small.num_anchors() else: anchor_params = None num_anchors = None # Keras recommends initialising a multi-gpu model on the CPU to ease weight sharing, and to prevent OOM errors. # optionally wrap in a parallel model if multi_gpu > 1: from keras.utils import multi_gpu_model with tf.device('/cpu:0'): model = model_with_weights(backbone_retinanet( num_classes, num_anchors=num_anchors, modifier=modifier, inputs=inputs, distance=distance), weights=weights, skip_mismatch=True, config=copy.deepcopy(cfg), num_classes=num_classes) training_model = multi_gpu_model(model, gpus=multi_gpu) else: model = model_with_weights(backbone_retinanet(num_classes, num_anchors=num_anchors, modifier=modifier, inputs=inputs, distance=distance, cfg=cfg), weights=weights, skip_mismatch=True, config=copy.deepcopy(cfg), num_classes=num_classes) training_model = model try: from keras.utils import plot_model # Write the keras model plot into a file plot_path = os.path.join(cfg.tb_logdir, cfg.model_name) makedirs(plot_path) plot_model(training_model, to_file=(os.path.join(plot_path, cfg.network) + '.png'), show_shapes=True) except Exception: # TODO: Catch the particular exceptions print(traceback.format_exc()) print(sys.exc_info()[2]) # make prediction model prediction_model = retinanet_bbox( model=model, anchor_params=anchor_params, score_thresh_train=cfg.score_thresh_train, class_specific_filter=cfg.class_specific_nms) # compile model if distance: training_model.compile(loss={ 'regression': losses.smooth_l1(), 'classification': losses.focal(), 'distance': losses.smooth_l1(alpha=distance_alpha) }, optimizer=keras.optimizers.adam(lr=lr, clipnorm=0.001)) else: training_model.compile(loss={ 'regression': losses.smooth_l1(), 'classification': losses.focal(), }, optimizer=keras.optimizers.adam(lr=lr, clipnorm=0.001)) return model, training_model, prediction_model