# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Unfreeze_Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
#------------------------------------------------------#
if True:
lr = 1e-3
Batch_size = 8
Init_Epoch = 0
Freeze_Epoch = 50
optimizer = optim.Adam(net.parameters(),lr)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.92)
if Use_Data_Loader:
train_dataset = YoloDataset(lines[:num_train], (Config["img_h"], Config["img_w"]), True)
val_dataset = YoloDataset(lines[num_train:], (Config["img_h"], Config["img_w"]), False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4,pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
else:
gen = Generator(Batch_size, lines[:num_train],
(Config["img_h"], Config["img_w"])).generate(True)
gen_val = Generator(Batch_size, lines[num_train:],
(Config["img_h"], Config["img_w"])).generate(False)
epoch_size = num_train//Batch_size
epoch_size_val = num_val//Batch_size
#------------------------------------#
# 冻结一定部分训练
lr = 1e-3
Batch_size = 4
Init_Epoch = 0
Freeze_Epoch = 50
#----------------------------------------------------------------------------#
# 我在实际测试时，发现optimizer的weight_decay起到了反作用，
# 所以去除掉了weight_decay，大家也可以开起来试试，一般是weight_decay=5e-4
#----------------------------------------------------------------------------#
optimizer = optim.Adam(net.parameters(),lr)
if Cosine_lr:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
else:
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.92)
train_dataset = YoloDataset(lines[:num_train], (input_shape[0], input_shape[1]), mosaic=mosaic, is_train=True)
val_dataset = YoloDataset(lines[num_train:], (input_shape[0], input_shape[1]), mosaic=False, is_train=False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4,pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
epoch_size = num_train // Batch_size
epoch_size_val = num_val // Batch_size
if epoch_size == 0 or epoch_size_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
#------------------------------------#
# 冻结一定部分训练
#------------------------------------#
for param in model.backbone.parameters():
anchors = get_anchors(anchors_path)
num_classes = len(class_names)
# 0.2用于验证，0.8用于训练
val_split = 0.2
with open(annotation_path) as f:
lines = f.readlines()
# 打乱数据顺序
np.random.shuffle(lines)
# 训练/验证集总数
num_val = int(len(lines) * val_split)
num_train = len(lines) - num_val
#-----------------------------------#
# 构建dataset
#-----------------------------------#
train_dataset = YoloDataset(lines[:num_train], input_shape, mosaic=mosaic)
'''
it = iter(train_dataset)
for i in range(24):
ii = next(it)
a, b = ii[0], ii[1]
print(i)
print(a.shape, b.shape)
'''
val_dataset = YoloDataset(lines[num_train:], input_shape, mosaic=False)
#-----------------------------------#
# 创建模型、加载预训练权重
#-----------------------------------#
model = YoloBody(3, num_classes)
# 加快模型训练的效率
# Freeze_Epoch为冻结训练的世代
# Unfreeze_Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
#------------------------------------------------------#
if False:
lr = 1e-3
Batch_size = 8
Init_Epoch = 0
Freeze_Epoch = 50
optimizer = optim.Adam(net.parameters(), lr)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=1,
gamma=0.92)
train_dataset = YoloDataset(lines[:num_train],
(input_shape[0], input_shape[1]), True)
val_dataset = YoloDataset(lines[num_train:],
(input_shape[0], input_shape[1]), False)
gen = DataLoader(train_dataset,
shuffle=True,
batch_size=Batch_size,
pin_memory=True,
drop_last=True,
collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset,
shuffle=True,
batch_size=Batch_size,
pin_memory=True,
drop_last=True,
collate_fn=yolo_dataset_collate)
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
optimizer = optim.Adam(model_train.parameters(), lr, weight_decay=5e-4)
if Cosine_lr:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=5,
eta_min=1e-5)
else:
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=1,
gamma=0.94)
train_dataset = YoloDataset(train_lines,
input_shape,
num_classes,
mosaic=mosaic,
train=True)
val_dataset = YoloDataset(val_lines,
input_shape,
num_classes,
mosaic=False,
train=False)
gen = DataLoader(train_dataset,
shuffle=True,
batch_size=batch_size,
num_workers=num_workers,
pin_memory=True,
drop_last=True,
collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset,
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
num_val = int(len(lines) * val_split)
num_train = len(lines) - num_val
# 优化器,最开始使用1e-3的学习率可以收敛的更快
lr = opt.lr * opt.gamma**(2 * opt.start_epoch) # 初始化学习率
print(lr)
optimizer = optim.Adam(net.parameters(), lr)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=1,
gamma=opt.gamma)
train_dataset = YoloDataset(lines[:num_train],
(Config["img_h"], Config["img_w"]))
val_dataset = YoloDataset(lines[num_train:],
(Config["img_h"], Config["img_w"]))
gen = DataLoader(train_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
drop_last=True,
collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
drop_last=True,
#---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
#---------------------------------------#
# 判断每一个世代的长度
#---------------------------------------#
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
#---------------------------------------#
# 构建数据集加载器。
#---------------------------------------#
train_dataset = YoloDataset(train_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, mosaic=mosaic, train = True)
val_dataset = YoloDataset(val_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, mosaic=False, train = False)
gen = DataLoader(train_dataset, shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset , shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
#---------------------------------------#
# 开始模型训练
#---------------------------------------#
for epoch in range(Init_Epoch, UnFreeze_Epoch):
#---------------------------------------#
# 如果模型有冻结学习部分
# 则解冻，并设置参数
#---------------------------------------#
if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Unfreeze_Epoch总训练世代
#------------------------------------------------------#
if True:
lr = 0.0008
Batch_size = 32
Init_Epoch = 0
Freeze_Epoch = 20
optimizer = optim.Adam(net.parameters(), lr)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=1,
gamma=0.92)
train_dataset = YoloDataset(train_lines,
(input_shape[0], input_shape[1]), True)
val_dataset = YoloDataset(val_lines, (input_shape[0], input_shape[1]),
False)
gen = DataLoader(train_dataset,
shuffle=True,
batch_size=Batch_size,
num_workers=12,
pin_memory=True,
drop_last=True,
collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset,
shuffle=True,
batch_size=Batch_size,
num_workers=12,
pin_memory=True,
drop_last=True,
# '''
lr = 1e-3
Batch_size = 4
Init_Epoch = 0
Freeze_Epoch = 25
Use_Data_Loader = True
optimizer = optim.Adam(model.parameters(), lr, weight_decay=5e-4)
if Cosine_lr:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-5)
else:
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
if Use_Data_Loader:
print('+' * 50)
train_dataset = YoloDataset(train_lines, (input_shape[0], input_shape[1]), mosaic=mosaic)
val_dataset = YoloDataset(val_lines, (input_shape[0], input_shape[1]), mosaic=False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=2, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=2, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
else:
gen = Generator(Batch_size, train_lines,
(input_shape[0], input_shape[1])).generate(mosaic=mosaic)
gen_val = Generator(Batch_size, val_lines,
(input_shape[0], input_shape[1])).generate(mosaic=False)
epoch_size = int(max(1, num_train // Batch_size // 2.5)) if mosaic else max(1, num_train // Batch_size)
epoch_size_val = num_val // Batch_size
for param in model.backbone.parameters():
param.requires_grad = False
