#------------------------------------------------------#
# 载入预训练权重
#------------------------------------------------------#
print('Load weights {}.'.format(model_path))
model.load_weights(model_path, by_name=True, skip_mismatch=True)
if focal_loss:
if dice_loss:
loss = dice_loss_with_Focal_Loss(cls_weights)
else:
loss = Focal_Loss(cls_weights)
else:
if dice_loss:
loss = dice_loss_with_CE(cls_weights)
else:
loss = CE(cls_weights)
#---------------------------#
# 读取数据集对应的txt
#---------------------------#
with open(
os.path.join(VOCdevkit_path,
"VOC2007/ImageSets/Segmentation/train.txt"),
"r") as f:
train_lines = f.readlines()
with open(
os.path.join(VOCdevkit_path,
"VOC2007/ImageSets/Segmentation/val.txt"), "r") as f:
val_lines = f.readlines()
num_train = len(train_lines)
num_val = len(val_lines)
# 打开数据集的txt
with open("VOCdevkit/VOC2007/ImageSets/Segmentation/val.txt", "r") as f:
val_lines = f.readlines()
if backbone == "mobilenet":
freeze_layers = 146
else:
freeze_layers = 172
for i in range(freeze_layers):
model.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(
freeze_layers, len(model.layers)))
loss = dice_loss_with_CE() if dice_loss else CE()
#------------------------------------------------------#
# 主干特征提取网络特征通用，冻结训练可以加快训练速度
# 也可以在训练初期防止权值被破坏。
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
#------------------------------------------------------#
if True:
Lr = 1e-4
Init_Epoch = 0
Freeze_Epoch = 50
Batch_size = 8
if Use_Data_Loader:
# 权值文件位置
#-------------------------------------------#
model_path = "model_data/pspnet_mobilenetv2.h5"
model.load_weights(model_path, by_name=True, skip_mismatch=True)
# 打开训练集
with open(os.path.join(dataset_path, "ImageSets/Segmentation/train.txt"),
"r") as f:
train_lines = f.readlines()
# 打开验证集
with open(os.path.join(dataset_path, "ImageSets/Segmentation/val.txt"),
"r") as f:
val_lines = f.readlines()
loss = dice_loss_with_CE() if dice_loss else CE()
#-------------------------------------------------------------------------------#
# 训练参数的设置
# logging表示tensorboard的保存地址
# checkpoint用于设置权值保存的细节，period用于修改多少epoch保存一次
# reduce_lr用于设置学习率下降的方式
# early_stopping用于设定早停，val_loss多次不下降自动结束训练，patience:多少轮不发生变化才停止
#-------------------------------------------------------------------------------#
checkpoint_period = ModelCheckpoint(
log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
monitor='val_loss',
save_weights_only=True,
save_best_only=False,
period=1)
reduce_lr = ReduceLROnPlateau(monitor='val_loss',
factor=0.5,
def main():
args = get_arguments()
h, w, c = map(int, args.input_size.split(','))
input_size = (h, w, c)
model = pspnet(args.num_classes, input_size, downsample_factor=downsample_factor,backbone=BACKBONE, aux_branch=aux_branch)
model.summary()
model.load_weights(args.model_path, by_name=True, skip_mismatch=True)
tf.set_random_seed(args.random_seed)
with open('list/cityscapes_train_list.txt', 'r') as f :
train_lines = f.readlines()
with open('list/cityscapes_val_list.txt', 'r') as f:
val_lines = f.readlines()
checkpoint_period = ModelCheckpoint(args.log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
monitor='val_loss',save_weights_only=True, save_best_only=False,period=1)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1)
early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)
tensorboard = TensorBoard(log_dir=args.log_dir)
for i in range(RESNET_FREEZE):
model.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(RESNET_FREEZE, len(model.layers),args.batch_size))
if True:
lr = 1e-4
Init_Epoch = 0
Freeze_Epoch = 50
BATCH_SIZE = args.batch_size
# model.compile(loss=dice_loss_with_CE() if dic_loss else CE(),optimizer=Adam(lr=lr),metrics=[f_score()])
model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=lr), metrics=['accuracy'])
gen = Generator(args.data_dir,args.data_list,BATCH_SIZE,input_size[:2],args.ignore_label,IMG_MEAN,aux_branch,len(train_lines),args.num_classes).generate()
gen_val = Generator(args.data_dir,args.data_val,BATCH_SIZE,input_size[:2],args.ignore_label,IMG_MEAN,aux_branch,len(val_lines),args.num_classes).generate(False)
model.fit_generator(gen,
steps_per_epoch=max(1, len(train_lines)//BATCH_SIZE)//4,
validation_data=gen_val,
validation_steps=max(1, len(val_lines)//BATCH_SIZE)//4,
epochs=Freeze_Epoch,
initial_epoch=Init_Epoch,
callbacks=[checkpoint_period, reduce_lr, tensorboard])
for i in range(RESNET_FREEZE):
model.layers[i].trainable = True
if True:
lr = 1e-5
Freeze_Epoch = 50
Unfreeze_Epoch = 100
BATCH_SIZE = args.batch_size / 2
model.compile(loss=dice_loss_with_CE() if dic_loss else CE(),
optimizer=Adam(lr=lr),
metrics=[f_score()])
print(
'Freeze the first {} layers of total {} layers.'.format(RESNET_FREEZE, len(model.layers), args.batch_size))
gen = Generator(args.data_dir,args.data_list,BATCH_SIZE,input_size[:2],args.ignore_label,IMG_MEAN,aux_branch,len(train_lines),args.num_classes).generate()
gen_val = Generator(args.data_dir,args.data_val,BATCH_SIZE,input_size[:2],args.ignore_label,IMG_MEAN,aux_branch,len(val_lines),args.num_classes).generate(False)
model.fit_generator(gen,
steps_per_epoch=max(1, len(train_lines) // BATCH_SIZE),
validation_data=gen_val,
validation_steps=max(1, len(val_lines) // BATCH_SIZE),
epochs=Unfreeze_Epoch,
initial_epoch=Freeze_Epoch,
callbacks=[checkpoint_period, reduce_lr, tensorboard])