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())