def load_model(filepath, backbone_name='resnet50', convert=False, nms=False): """ Loads a retinanet model using the correct custom objects. # Arguments filepath: one of the following: - string, path to the saved model, or - h5py.File object from which to load the model backbone_name: Backbone with which the model was trained. convert: Boolean, whether to convert the model to an inference model. nms: Boolean, whether to add NMS filtering to the converted model. Only valid if convert=True. # Returns A keras.models.Model object. # Raises ImportError: if h5py is not available. ValueError: In case of an invalid savefile. """ import keras.models model = keras.models.load_model( filepath, custom_objects=backbone(backbone_name).custom_objects) if convert: from object_detector_retinanet.keras_retinanet.models.retinanet import retinanet_bbox model = retinanet_bbox(model=model, nms=nms) return model
def create_models(backbone_retinanet, num_classes, weights, multi_gpu=0, freeze_backbone=False): """ Creates three models (model, training_model, prediction_model). Args backbone_retinanet : A function to call to create a retinanet model with a given backbone. num_classes : The number of classes to train. weights : The weights to load into the model. multi_gpu : The number of GPUs to use for training. freeze_backbone : If True, disables learning for the backbone. Returns model : The base model. This is also the model that is saved in snapshots. training_model : The training model. If multi_gpu=0, this is identical to model. prediction_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 # 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 num_anchors = AnchorParameters.default.num_anchors() submodels = iou_submodels(num_classes=num_classes, num_anchors=num_anchors) if multi_gpu > 1: with tf.device('/cpu:0'): model = model_with_weights(backbone_retinanet(num_classes, modifier=modifier, submodels=submodels), weights=weights, skip_mismatch=True) iou_model = multi_gpu_model(model, gpus=multi_gpu) else: model = model_with_weights(backbone_retinanet(num_classes, modifier=modifier, submodels=submodels), weights=weights, skip_mismatch=True) # make prediction model iou_model = retinanet_iou(model=model) prediction_model = retinanet_bbox(model=model) # compile model iou_model.compile(loss={'iou_regression_output': losses.iou_score()}, optimizer=keras.optimizers.adam(lr=1e-5, clipnorm=0.001) # optimizer=keras.optimizers.SGD() ) return model, iou_model, prediction_model
def main(args=None): # parse arguments if args is None: args = sys.argv[1:] args = parse_args(args) if DEBUG_MODE: args.steps = 10 # create object that stores backbone information backbone = models.backbone(args.backbone) # make sure keras is the minimum required version check_keras_version() # optionally choose specific GPU use_cpu = False if args.gpu: gpu_num = args.gpu else: gpu_num = str(0) if use_cpu: os.environ["CUDA_VISIBLE_DEVICES"] = str(666) else: os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num keras.backend.tensorflow_backend.set_session(get_session()) # Weights and logs saves in a new locations stmp = time.strftime("%c").replace(" ", "_") args.snapshot_path = os.path.join(args.snapshot_path, stmp) args.tensorboard_dir = os.path.join(args.tensorboard_dir, stmp) print("Weights will be saved in {}".format(args.snapshot_path)) print("Logs will be saved in {}".format(args.tensorboard_dir)) create_folder(args.snapshot_path) create_folder(args.tensorboard_dir) # create the generators train_generator, validation_generator = create_generators(args) print('train_size:{},val_size:{}'.format(train_generator.size(), validation_generator.size())) # create the model if args.snapshot is not None: print('Loading model, this may take a second...') model = models.load_model(args.snapshot, backbone_name=args.backbone) training_model = model prediction_model = retinanet_bbox(model=model) else: weights = os.path.join(os.path.join(root_dir(), args.weights)) # default to imagenet if nothing else is specified if weights is None and args.imagenet_weights: weights = backbone.download_imagenet() 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=args.multi_gpu, freeze_backbone=args.freeze_backbone) # print model summary # print(model.summary()) # this lets the generator compute backbone layer shapes using the actual backbone model if 'vgg' in args.backbone or 'densenet' in args.backbone: compute_anchor_targets = functools.partial( anchor_targets_bbox, shapes_callback=make_shapes_callback(model)) train_generator.compute_anchor_targets = compute_anchor_targets if validation_generator is not None: validation_generator.compute_anchor_targets = compute_anchor_targets # create the callbacks callbacks = create_callbacks( model, training_model, prediction_model, validation_generator, args, ) # start training training_model.fit_generator(generator=train_generator, steps_per_epoch=args.steps, epochs=args.epochs, verbose=1, callbacks=callbacks, validation_data=validation_generator, validation_steps=validation_generator.size())