def load_model(filepath,
backbone_name='resnet50',
convert=False,
nms=True,
class_specific_filter=True):
""" 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.
class_specific_filter : Whether to use class specific filtering or filter for the best scoring class only.
# Returns
A keras.models.Model object.
# Raises
ImportError: if h5py is not available.
ValueError: In case of an invalid savefile.
"""
model = keras.models.load_model(
filepath, custom_objects=backbone(backbone_name).custom_objects)
if convert:
from models.retinanet import retinanet_bbox
print("Starting to convert model...")
model = retinanet_bbox(model=model,
nms=nms,
class_specific_filter=class_specific_filter)
return model
def create_models(backbone_retinanet,
num_classes,
weights,
args,
num_gpus=0,
freeze_backbone=False,
lr=1e-5,
config=None):
"""
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.
num_gpus: The number of GPUs to use for training.
freeze_backbone: If True, disables learning for the backbone.
config: Config parameters, None indicates the default configuration.
Returns
model: The base model. This is also the model that is saved in snapshots.
training_model: The training model. If num_gpus=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
if num_gpus > 1:
from keras.utils import multi_gpu_model
with tf.device('/cpu:0'):
model = model_with_weights(backbone_retinanet(num_classes,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = multi_gpu_model(model, gpus=num_gpus)
else:
model = model_with_weights(backbone_retinanet(num_classes,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = model
# make prediction model
prediction_model = retinanet_bbox(model=model)
# compile model
training_model.compile(loss={
'regression':
losses.iou_loss(args.loss, args.loss_weight),
'classification':
losses.focal(),
'centerness':
losses.bce(),
},
optimizer=keras.optimizers.adam(lr=lr))
return model, training_model, prediction_model
def train(args):
# create object that stores backbone information
backbone = model_backbone.backbone(args.backbone)
# optionally choose specific GPU
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
keras.backend.tensorflow_backend.set_session(get_session())
# create the generators
print("Going to get the training and validation generators...")
train_generator, validation_generator = create_generators(
args, backbone.preprocess_image)
# create the model
if args.snapshot is not None:
print('Loading model: {} \nThis may take a second...'.format(
args.snapshot))
model = model_backbone.load_model(args.snapshot,
backbone_name=args.backbone)
training_model = model
prediction_model = retinanet_bbox(model=model)
else:
weights = 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())
# create the callbacks
callbacks = create_callbacks(
model,
training_model,
prediction_model,
validation_generator,
args,
)
# start training
print("Started training...")
training_model.fit_generator(
generator=train_generator,
steps_per_epoch=args.steps,
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
)
print("Training Complete")
def create_models(backbone_retinanet,
num_classes,
weights,
multi_gpu=1,
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
if multi_gpu > 1:
with tf.device('/cpu:0'):
model = model_with_weights(backbone_retinanet(num_classes,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = multi_gpu_model(model, gpus=multi_gpu)
else:
model = model_with_weights(backbone_retinanet(num_classes,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = model
# make prediction model
prediction_model = retinanet_bbox(model=model)
# compile model
training_model.compile(loss={
'regression': losses.smooth_l1(),
'classification': losses.focal()
},
optimizer=keras.optimizers.adam(lr=1e-7,
clipnorm=0.001))
return model, training_model, prediction_model
def _create_models(self,
backbone_retinanet,
num_classes,
weights,
freeze_backbone=False,
lr=1e-5):
"""
Creates three models (model, training_model, prediction_model).
Parameters
----------
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.
config : Config parameters, None indicates the default configuration.
Returns
-------
model : The base model.
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
anchor_params = None
num_anchors = None
model = self._model_with_weights(backbone_retinanet(
num_classes, num_anchors=num_anchors, modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = model
prediction_model = retinanet_bbox(model=model,
anchor_params=anchor_params)
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
def main(config_file=None):
cwd = os.getcwd()
# parse configuration file
if config_file is None:
config_file = sys.argv[-1]
config_file = os.path.join(cwd, config_file)
config_file_name = config_file.split('/')[-1]
configs = parse_config(config_file)
# save config file
if configs['Train']['save_configs']:
# confirm save dir
config_save_path = 'logs/' + configs['Name']
config_save_path = os.path.join(cwd, config_save_path)
if not os.path.exists(config_save_path):
os.mkdir(os.path.join(cwd, config_save_path))
# copy config file
config_dst_name = configs['Name'] + '.json'
config_file_dst = os.path.join(config_save_path, config_dst_name)
shutil.copy(config_file, config_file_dst)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
if configs['Train']['gpu']:
os.environ['CUDA_VISIBLE_DEVICES'] = configs['Train']['gpu']
keras.backend.tensorflow_backend.set_session(get_session())
# create the generators
train_generator, validation_generator = create_generators(configs)
# create the model
if configs['Train']['load_snapshot'] is not None:
print('Loading model, this may take a second...')
model = models.load_model(configs['Train']['load_snapshot'],
backbone=configs['Train']['backbone'])
training_model = prediction_model = model
else:
weights = configs['Train']['weights']
# default to imagenet if nothing else is specified
if weights is None and configs['Train']['imagenet_weights']:
weights = models.download_imagenet(configs['Train']['backbone'])
print('Creating model, this may take a second...')
backbone = configs['Train']['backbone']
num_classes = train_generator.num_classes()
multi_gpu = configs['Train']['multi_gpu']
freeze_backbone = configs['Train']['freeze_backbone']
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
if multi_gpu > 1:
with tf.device('/cpu:0'):
retinanet = models.retinanet_backbone(
configs['Train']['backbone'])(num_classes,
backbone=backbone,
modifier=modifier)
model = model_with_weights(retinanet,
weights=weights,
skip_mismatch=True)
training_model = multi_gpu_model(model, gpus=multi_gpu)
else:
retinanet = models.retinanet_backbone(
configs['Train']['backbone'])(num_classes,
backbone=backbone,
modifier=modifier)
training_model = model = model_with_weights(retinanet,
weights=weights,
skip_mismatch=True)
# make prediction model
prediction_model = retinanet_bbox(model=model,
anchor_param=configs['Anchors'])
# compile model
subnet_loss = {
'regression': losses.smooth_l1(),
'classification': losses.focal()
}
# run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
# run_metadata = tf.RunMetadata()
if configs['Train']['lr_multiplier_layer']:
optimizer = AdamWithLRMult(
lr=configs['Train']['init_lr'],
lr_multipliers=configs['Train']['lr_multiplier_layer'],
debug_verbose=False,
clipnorm=0.001)
else:
optimizer = keras.optimizers.adam(configs['Train']['init_lr'],
clipnorm=0.001)
training_model.compile(loss=subnet_loss, optimizer=optimizer)
# training_model.compile(loss=subnet_loss, optimizer=keras.optimizers.adam(configs['Train']['init_lr'], clipnorm=0.001))
# this lets the generator compute backbone layer shapes using the actual backbone model
if 'vgg' in configs['Train']['backbone'] or 'densenet' in configs['Train'][
'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,
configs,
)
# start training
training_model.fit_generator(
train_generator,
validation_data=validation_generator,
validation_steps=39,
steps_per_epoch=configs['Train']['steps'],
epochs=configs['Train']['epochs'],
verbose=1,
callbacks=callbacks,
)
def main(args=None):
# create object that stores backbone information
backbone = models.backbone(args.backbone)
# create the generators
train_generator, validation_generator = create_generators(
args, backbone.preprocess_image)
# 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
anchor_params = None
if args.config and 'anchor_parameters' in args.config:
anchor_params = parse_anchor_parameters(args.config)
prediction_model = retinanet_bbox(model=model,
anchor_params=anchor_params)
# compile model
training_model.compile(
loss={
'regression': losses.iou_loss(args.loss, args.loss_weight),
'classification': losses.focal(),
'centerness': losses.bce(),
},
optimizer=keras.optimizers.Adam(lr=1e-5)
# optimizer=keras.optimizers.sgd(lr=1e-5, momentum=0.9, decay=1e-5, nesterov=True)
)
else:
weights = 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,
num_gpus=args.num_gpus,
freeze_backbone=args.freeze_backbone,
lr=args.lr,
config=args.config,
args=args)
parallel_model = multi_gpu_model(training_model, gpus=2)
parallel_model.compile(loss={
'regression':
losses.iou_loss(args.loss, args.loss_weight),
'classification':
losses.focal(),
'centerness':
losses.bce(),
},
optimizer=keras.optimizers.Adam(lr=1e-4))
# 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:
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,
training_model,
prediction_model,
validation_generator,
args,
)
if not args.compute_val_loss:
validation_generator = None
# start training
parallel_model.fit_generator(generator=train_generator,
steps_per_epoch=len(train_generator),
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
validation_data=validation_generator)
def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
# 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
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
keras.backend.tensorflow_backend.set_session(get_session())
# optionally load config parameters
if args.config:
args.config = read_config_file(args.config)
# create the generators
train_generator, validation_generator = create_generators(
args, backbone.preprocess_image)
# 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
anchor_params = None
if args.config and 'anchor_parameters' in args.config:
anchor_params = parse_anchor_parameters(args.config)
prediction_model = retinanet_bbox(model=model,
anchor_params=anchor_params)
# compile model
training_model.compile(
loss={
'regression': losses.iou(),
'classification': losses.focal(),
'centerness': losses.bce(),
},
optimizer=keras.optimizers.adam(lr=1e-5)
# optimizer=keras.optimizers.sgd(lr=1e-5, momentum=0.9, decay=1e-5, nesterov=True)
)
else:
weights = 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,
num_gpus=args.num_gpus,
freeze_backbone=args.freeze_backbone,
lr=args.lr,
config=args.config)
# 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:
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,
training_model,
prediction_model,
validation_generator,
args,
)
if not args.compute_val_loss:
validation_generator = None
# start training
return training_model.fit_generator(
generator=train_generator,
initial_epoch=0,
steps_per_epoch=args.steps,
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
workers=args.workers,
use_multiprocessing=args.multiprocessing,
max_queue_size=args.max_queue_size,
validation_data=validation_generator)
def create_models(fl_gamma,
fl_alpha,
r_weight,
c_weight,
p_weight,
train_type,
sample_t,
backbone_retinanet,
num_classes,
weights,
class_weights,
loss_weights,
multi_gpu=0,
freeze_backbone=False,
lr=1e-5,
config=None):
""" 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.
config : Config parameters, None indicates the default configuration.
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
# load anchor parameters, or pass None (so that defaults will be used)
anchor_params = None
num_anchors = None
if config and 'anchor_parameters' in config:
anchor_params = parse_anchor_parameters(config)
num_anchors = anchor_params.num_anchors()
# 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),
weights=weights,
skip_mismatch=True)
training_model = multi_gpu_model(model, gpus=multi_gpu)
else:
model = model_with_weights(backbone_retinanet(num_classes,
train_type,
num_anchors=num_anchors,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = model
# make prediction model
prediction_model = retinanet_bbox(model=model, anchor_params=anchor_params)
if train_type == "single":
# compile model
training_model.compile(
loss={
'regression':
losses.smooth_l1(r_weight),
## HERE THE INPUT ARGUMENTS CAN BE GIVEN: default focal(alpha=0.25, gamma=2.0)
## gamma: the actual "focusing parameter"
'classification':
losses.focal(c_weight, fl_alpha, fl_gamma, class_weights)
},
optimizer=keras.optimizers.adam(lr=lr, clipnorm=0.001),
#sample_weight_mode="temporal",
#sample_weights=sample_t
)
elif train_type == "multi":
training_model.compile(
loss={
'regression':
losses.smooth_l1(r_weight),
## HERE THE INPUT ARGUMENTS CAN BE GIVEN: default focal(alpha=0.25, gamma=2.0)
## gamma: the actual "focusing parameter"
'classification_artefact':
losses.focal(c_weight, fl_alpha, fl_gamma),
'classification_polyp':
losses.focal2(p_weight, fl_alpha, fl_gamma)
},
optimizer=keras.optimizers.adam(lr=lr, clipnorm=0.001)
#loss_weights=loss_weights
)
return model, training_model, prediction_model
def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
# create object that stores backbone information
#backbone = models.backbone(args.backbone)
backbone = models.backbone(args.backbone)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
keras.backend.tensorflow_backend.set_session(get_session())
# optionally load config parameters
if args.config:
args.config = read_config_file(args.config)
# create the generators
#print(args)
train_generator, validation_generator = create_generators(
args, backbone.preprocess_image)
# Log configs
#run.log('batch-size', args.batch_size)
#run.log('gamma', args.fl_gamma)
#run.log('alpha', args.fl_alpha)
#run.log('lr', args.lr)
#run.log('neg-overlap', args.neg_overlap)
#run.log('pos-overlap', args.pos_overlap)
#run.log('fpn-layers', args.fpn_layers)
if args.class_weights is not None:
if args.class_weights == "cw1":
polyp_weight = 0.25
#class_weights = {'classification': {0:a_w, 1:a_w, 2:a_w, 3:a_w, 4:a_w, 5:a_w, 6:a_w, 7:polyp_weight}, 'regression': {0:0.25, 1:0.25, 2:0.25, 3:0.25}}
#class_weights = {'classification': [polyp_weight, a_w, a_w, a_w, a_w, a_w, a_w, a_w]}
elif args.class_weights == "cw2":
polyp_weight = 0.5
elif args.class_weights == "cw3":
polyp_weight = 0.75
a_w = (1 - polyp_weight) / 7
#class_weights = {'classification': [polyp_weight, a_w, a_w, a_w, a_w, a_w, a_w, a_w]}
class_weights = [polyp_weight, a_w, a_w, a_w, a_w, a_w, a_w, a_w]
else:
class_weights = None
if args.loss_weights is None:
loss_weights = [1, 1]
elif args.loss_weights == "lw0":
loss_weights = [1, 1, 1]
elif args.loss_weights == "lw1":
loss_weights = [1, 1, 3]
elif args.loss_weights == "lw2":
loss_weights = [1, 1, 10]
elif args.loss_weights == "lw3":
loss_weights = [1, 1, 20]
# create the model
if args.snapshot is not None:
print('Loading model, this may take a second...')
model = models.load_model(os.path.join(args.data_dir, args.snapshot),
backbone_name=args.backbone)
training_model = model
anchor_params = None
if args.config and 'anchor_parameters' in args.config:
anchor_params = parse_anchor_parameters(args.config)
prediction_model = retinanet_bbox(model=model,
anchor_params=anchor_params)
else:
if args.weights is None and args.imagenet_weights:
weights = backbone.download_imagenet()
else:
weights = args.weights
# default to imagenet if nothing else is specified
## SO the file that is downloaded is actually only the weights
## this means that I should be able to use --weights to give it my own model
sample_test = np.array([[0.25, 0.25, 0.25, 0.25, 0, 0, 0, 0],
[10, 10, 10, 10, 10, 10, 10, 10]])
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,
class_weights=class_weights,
loss_weights=loss_weights,
multi_gpu=args.multi_gpu,
freeze_backbone=args.freeze_backbone,
lr=args.lr,
config=args.config,
fl_gamma=args.fl_gamma,
fl_alpha=args.fl_alpha,
c_weight=args.c_weight,
r_weight=args.r_weight,
p_weight=args.p_weight,
train_type=args.train_type,
sample_t=sample_test)
# 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:
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,
training_model,
prediction_model,
validation_generator,
args,
)
# Use multiprocessing if workers > 0
if args.workers > 0:
use_multiprocessing = True
else:
use_multiprocessing = False
temp_df = pd.read_csv(
os.path.join(args.data_dir, args.annotations),
names=["image_path", "x1", "y1", "x2", "y2", "object_id"])
im_count = len(set(list(temp_df.image_path)))
# start training
training_model.fit_generator(generator=train_generator,
steps_per_epoch=int(im_count /
args.batch_size),
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
workers=args.workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=args.max_queue_size
#class_weight=class_weights
)
def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
# optionally load config parameters
if args.config:
args.config = read_config_file(args, 'training')
# print("----------------------------------")
# print("ARGUMENTS IN CONFIG FILE:")
# for sec in args.config.sections():
# print(sec, "=", dict(args.config.items(sec)))
# print("----------------------------------")
#
# for arg in vars(args):
# print(arg, "=", getattr(args, arg))
# exit()
# 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
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
keras.backend.tensorflow_backend.set_session(get_session())
# create the generators
train_generator, validation_generator = create_generators(
args, backbone.preprocess_image)
# # Debuging
# for i in range(1):
# inputs, targets = train_generator.__getitem__(i)
# exit()
# 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)
# When using as a second step for fine-tuning
for layer in model.layers:
layer.trainable = True
training_model = model
anchor_params = None
if args.config and 'anchor_parameters' in args.config:
anchor_params = parse_anchor_parameters(args.config)
prediction_model = retinanet_bbox(model=model,
anchor_params=anchor_params)
###################################################################################### BRUNO
# compile model
training_model.compile(loss={
'regression': losses.smooth_l1(),
'classification': losses.focal()
},
optimizer=keras.optimizers.adam(lr=args.lr,
clipnorm=0.001))
else:
weights = 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,
lr=args.lr,
config=args.config)
# Print model design
# print(model.summary())
# print(training_model.summary())
# plot_model(model, to_file='model_plot.png', show_shapes=True, show_layer_names=True)
# exit()
# Get the number of samples in the training and validations datasets.
train_size = train_generator.size()
val_size = validation_generator.size()
print('Train on {} samples, val on {} samples, with batch size {}.'.format(
train_size, val_size, args.batch_size))
# this lets the generator compute backbone layer shapes using the actual backbone model
if 'vgg' in args.backbone or 'densenet' in args.backbone:
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,
training_model,
prediction_model,
validation_generator,
train_size,
[1e-6, 1e-4],
args,
)
# Use multiprocessing if workers > 0
if args.workers > 0:
use_multiprocessing = True
else:
use_multiprocessing = False
# check to see if we are attempting to find an optimal learning rate
# before training for the full number of epochs
if args.find_lr:
# initialize the learning rate finder and then train with learning
# rates ranging from 1e-10 to 1e+1
print("[INFO] Finding learning rate...")
lrf = LearningRateFinder(training_model)
lrf.find(train_generator,
1e-10,
1e+1,
stepsPerEpoch=np.ceil((train_size / float(args.batch_size))),
batchSize=args.batch_size)
# plot the loss for the various learning rates and save the
# resulting plot to disk
lrf.plot_loss()
plt.savefig("lrfind_plot.png")
# save values into a csv file
lrf.save_csv("lr_loss.csv")
# gracefully exit the script so we can adjust our learning rates
# in the config and then train the network for our full set of
# epochs
print("[INFO] Learning rate finder complete")
print("[INFO] Examine plot and adjust learning rates before training")
sys.exit(0)
# Number of epochs and steps for training new layers
n_epochs = 350
steps = train_size // args.batch_size
# start training
training_model.fit_generator(generator=train_generator,
validation_data=validation_generator,
steps_per_epoch=steps,
epochs=n_epochs,
verbose=1,
callbacks=callbacks,
workers=args.workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=args.max_queue_size)
# Unfreeze and continue training, to fine-tune.
# Train longer if the result is not good.
if True:
# for layer in model.layers:
# if layer.name is 'bn_conv1':
# print("Before\t-> Trainable: {}, Freeze: {}".format(layer.trainable, layer.freeze))
for layer in model.layers:
layer.trainable = True
# recompile to apply the change
model.compile(
loss={
'regression': losses.smooth_l1(),
'classification': losses.focal()
},
# Learning rate must be lower for training the entire network
optimizer=keras.optimizers.adam(lr=args.lr * 0.1, clipnorm=0.001),
metrics=['accuracy'])
if args.multi_gpu > 1:
from keras.utils import multi_gpu_model
with tf.device('/gpu:1'):
training_model = multi_gpu_model(model, gpus=args.multi_gpu)
else:
training_model = model
# recompile to apply the change
training_model.compile(
loss={
'regression': losses.smooth_l1(),
'classification': losses.focal()
},
# Learning rate must be lower for training the entire network
optimizer=keras.optimizers.adam(lr=args.lr * 0.1, clipnorm=0.001),
metrics=['accuracy'])
print('Unfreezing all layers.')
# for layer in model.layers:
# if layer.name is 'bn_conv1':
# print("After\t-> Trainable: {}, Freeze: {}".format(layer.trainable, layer.freeze))
# Print training_model design
# print(model.summary())
# print(training_model.summary())
# create the callbacks
callbacks = create_callbacks(
model,
training_model,
prediction_model,
validation_generator,
train_size,
[1e-8, 1e-6],
args,
)
batch_size = 2 # note that more GPU memory is required after unfreezing the body
steps = train_size // batch_size
print('Train on {} samples, val on {} samples, with batch size {}.'.
format(train_size, val_size, batch_size))
training_model.fit_generator(generator=train_generator,
validation_data=validation_generator,
steps_per_epoch=steps,
epochs=args.epochs,
initial_epoch=n_epochs,
verbose=1,
callbacks=callbacks,
workers=args.workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=args.max_queue_size)
def create_models(backbone_retinanet,
num_classes,
weights,
multi_gpu=2,
freeze_backbone=False,
lr=1e-5,
config=None):
"""!@brief
Creates three models (model, training_model, prediction_model).
@param backbone_retinanet : A function to call to create a retinanet model with a given backbone.
@param num_classes : The number of classes to train.
@param weights : The weights to load into the model.
@param multi_gpu : The number of GPUs to use for training.
@param freeze_backbone : If True, disables learning for the backbone.
@param config : Config parameters, None indicates the default configuration.
@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
# load anchor parameters, or pass None (so that defaults will be used)
anchor_params = None
num_anchors = None
if config and 'anchor_parameters' in config:
anchor_params = parse_anchor_parameters(config)
num_anchors = anchor_params.num_anchors()
# 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('/gpu:1'):
model = model_with_weights(backbone_retinanet(
num_classes, num_anchors=num_anchors, modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = multi_gpu_model(
model,
gpus=multi_gpu) # , cpu_merge=False) # <- recommended for NV-Link
else:
model = model_with_weights(backbone_retinanet(num_classes,
num_anchors=num_anchors,
modifier=modifier),
weights=weights,
skip_mismatch=True)
training_model = model
# make prediction model
prediction_model = retinanet_bbox(model=model, anchor_params=anchor_params)
# compile model
training_model.compile(loss={
'regression': losses.smooth_l1(),
'classification': losses.focal()
},
optimizer=keras.optimizers.adam(lr=lr,
clipnorm=0.001),
metrics=['accuracy'])
return model, training_model, prediction_model