def get_yolo3_train_model(model_type, anchors, num_classes, weights_path=None, freeze_level=1, optimizer=Adam(lr=1e-3, decay=0), label_smoothing=0, elim_grid_sense=False, model_pruning=False, pruning_end_step=10000):
'''create the training model, for YOLOv3'''
#K.clear_session() # get a new session
num_anchors = len(anchors)
#YOLOv3 model has 9 anchors and 3 feature layers but
#Tiny YOLOv3 model has 6 anchors and 2 feature layers,
#so we can calculate feature layers number to get model type
num_feature_layers = num_anchors//3
#feature map target value, so its shape should be like:
# [
# (image_height/32, image_width/32, 3, num_classes+5),
# (image_height/16, image_width/16, 3, num_classes+5),
# (image_height/8, image_width/8, 3, num_classes+5)
# ]
y_true = [Input(shape=(None, None, 3, num_classes+5), name='y_true_{}'.format(l)) for l in range(num_feature_layers)]
model_body, backbone_len = get_yolo3_model(model_type, num_feature_layers, num_anchors, num_classes, model_pruning=model_pruning, pruning_end_step=pruning_end_step)
print('Create {} {} model with {} anchors and {} classes.'.format('Tiny' if num_feature_layers==2 else '', model_type, num_anchors, num_classes))
print('model layer number:', len(model_body.layers))
if weights_path:
model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
print('Load weights {}.'.format(weights_path))
if freeze_level in [1, 2]:
# Freeze the backbone part or freeze all but final feature map & input layers.
num = (backbone_len, len(model_body.layers)-3)[freeze_level-1]
for i in range(num): model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
elif freeze_level == 0:
# Unfreeze all layers.
for i in range(len(model_body.layers)):
model_body.layers[i].trainable= True
print('Unfreeze all of the layers.')
model_loss, location_loss, confidence_loss, class_loss = Lambda(yolo3_loss, name='yolo_loss',
arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5, 'label_smoothing': label_smoothing, 'elim_grid_sense': elim_grid_sense})(
[*model_body.output, *y_true])
model = Model([model_body.input, *y_true], model_loss)
loss_dict = {'location_loss':location_loss, 'confidence_loss':confidence_loss, 'class_loss':class_loss}
add_metrics(model, loss_dict)
model.compile(optimizer=optimizer, loss={
# use custom yolo_loss Lambda layer.
'yolo_loss': lambda y_true, y_pred: y_pred})
return model
def get_yolo2_train_model(model_type, anchors, num_classes, weights_path=None, freeze_level=1, optimizer=Adam(lr=1e-3, decay=1e-6), label_smoothing=0, model_pruning=False, pruning_end_step=10000):
'''create the training model, for YOLOv2'''
#K.clear_session() # get a new session
num_anchors = len(anchors)
# y_true in form of relative x, y, w, h, objectness, class
y_true_input = Input(shape=(None, None, num_anchors, 6))
model_body, backbone_len = get_yolo2_model(model_type, num_anchors, num_classes, model_pruning=model_pruning, pruning_end_step=pruning_end_step)
print('Create YOLOv2 {} model with {} anchors and {} classes.'.format(model_type, num_anchors, num_classes))
print('model layer number:', len(model_body.layers))
if weights_path:
model_body.load_weights(weights_path, by_name=True)#, skip_mismatch=True)
print('Load weights {}.'.format(weights_path))
if freeze_level in [1, 2]:
# Freeze the backbone part or freeze all but final feature map & input layers.
num = (backbone_len, len(model_body.layers)-2)[freeze_level-1]
for i in range(num): model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
elif freeze_level == 0:
# Unfreeze all layers.
for i in range(len(model_body.layers)):
model_body.layers[i].trainable= True
print('Unfreeze all of the layers.')
model_loss, location_loss, confidence_loss, class_loss = Lambda(yolo2_loss, name='yolo_loss',
arguments={'anchors': anchors, 'num_classes': num_classes, 'label_smoothing': label_smoothing})(
[model_body.output, y_true_input])
model = Model([model_body.input, y_true_input], model_loss)
loss_dict = {'location_loss':location_loss, 'confidence_loss':confidence_loss, 'class_loss':class_loss}
add_metrics(model, loss_dict)
model.compile(optimizer=optimizer, loss={
# use custom yolo_loss Lambda layer.
'yolo_loss': lambda y_true, y_pred: y_pred})
return model
def get_yolo3_train_model(model_type,
anchors,
num_classes,
input_shape,
weights_path=None,
freeze_level=1,
optimizer=Adam(lr=1e-3, decay=0),
label_smoothing=0):
"""create the training model, for YOLOv3"""
num_anchors = len(anchors)
num_feature_layers = num_anchors // 3
h, w = input_shape
# y_true = [Input(shape=(None, None, 3, num_classes + 5), name='y_true_{}'.format(l)) for l in
# range(num_feature_layers)]
y_true = [
Input(shape=(h // {
0: 32,
1: 16,
2: 8
}[l], w // {
0: 32,
1: 16,
2: 8
}[l], 3, num_classes + 5),
name='y_true_{}'.format(l)) for l in range(num_feature_layers)
]
model_body, backbone_len = get_yolo3_model(model_type, num_feature_layers,
num_anchors, num_classes)
print('Create {} model with {} anchors and {} classes.'.format(
model_type, num_anchors, num_classes))
print('model layer number:', len(model_body.layers))
if weights_path:
model_body.load_weights(weights_path,
by_name=True) # , skip_mismatch=True)
print('Load weights {}.'.format(weights_path))
if freeze_level in [1, 2]:
# Freeze the backbone part or freeze all but final feature map & input layers.
num = (backbone_len, len(model_body.layers) - 3)[freeze_level - 1]
for i in range(num):
model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(
num, len(model_body.layers)))
elif freeze_level == 0:
# Unfreeze all layers.
for i in range(len(model_body.layers)):
model_body.layers[i].trainable = True
print('Unfreeze all of the layers.')
use_focal_obj_loss = False
use_focal_loss = False
use_diou_loss = False
use_softmax_loss = False
model_loss, location_loss, confidence_loss, class_loss = Lambda(
yolo3_loss,
output_shape=(1, ),
name='yolo_loss',
arguments={
'anchors': anchors,
'num_classes': num_classes,
'ignore_thresh': 0.5,
'label_smoothing': label_smoothing,
'use_focal_obj_loss': use_focal_obj_loss,
'use_focal_loss': use_focal_loss,
'use_diou_loss': use_diou_loss,
'use_softmax_loss': use_softmax_loss
})([*model_body.output, *y_true])
model = Model([model_body.input, *y_true], model_loss)
loss_dict = {
'location_loss': location_loss,
'confidence_loss': confidence_loss,
'class_loss': class_loss
}
add_metrics(model, loss_dict)
return model
