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
min_val_loss = math.inf
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, writer,
mosaic=False)
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],
(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,
