Batch_size = 8
Init_Epoch = 0
Freeze_Epoch = 25
optimizer = optim.Adam(net.parameters(),lr)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)
if Use_Data_Loader:
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=Batch_size, num_workers=4, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, 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()
gen_val = Generator(Batch_size, lines[num_train:],
(Config["img_h"], Config["img_w"])).generate()
epoch_size = num_train//Batch_size
epoch_size_val = num_val//Batch_size
#------------------------------------#
# 冻结一定部分训练
#------------------------------------#
for param in model.backbone.parameters():
param.requires_grad = False
for epoch in range(Init_Epoch,Freeze_Epoch):
fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
lr_scheduler.step()
optimizer = optim.Adam(net.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.95)
if Use_Data_Loader:
train_dataset = YoloDataset(lines[:num_train], (input_shape[0], input_shape[1]), mosaic=mosaic)
val_dataset = YoloDataset(lines[num_train:], (input_shape[0], input_shape[1]), mosaic=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],
(input_shape[0], input_shape[1])).generate(mosaic = mosaic)
gen_val = Generator(Batch_size, lines[num_train:],
(input_shape[0], input_shape[1])).generate(mosaic = False)
epoch_size = max(1, num_train//Batch_size)
epoch_size_val = num_val//Batch_size
#------------------------------------#
# 冻结一定部分训练
#------------------------------------#
for param in model.backbone.parameters():
param.requires_grad = False
for epoch in range(Init_Epoch,Freeze_Epoch):
fit_one_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
lr_scheduler.step()
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
best_loss = 99999999.0
best_model_weights = copy.deepcopy(model.state_dict())
for epoch in range(Init_Epoch, Freeze_Epoch):
total_loss, val_loss = fit_one_epoch(model, yolo_losses, epoch, epoch_size, epoch_size_val, gen, gen_val,
Freeze_Epoch, Cuda, optimizer, lr_scheduler)
if total_loss < best_loss:
best_loss = total_loss