未验证 提交 308ac92c 编写于 作者: B Bubbliiiing 提交者: GitHub

Update train.py

上级 d67b456f
......@@ -160,9 +160,9 @@ if __name__ == "__main__":
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=8, pin_memory=True,
gen = DataLoader(train_dataset, batch_size=Batch_size, num_workers=0, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, batch_size=Batch_size, num_workers=8,pin_memory=True,
gen_val = DataLoader(val_dataset, batch_size=Batch_size, num_workers=0,pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
else:
gen = Generator(Batch_size, lines[:num_train],
......@@ -193,9 +193,9 @@ if __name__ == "__main__":
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=8, pin_memory=True,
gen = DataLoader(train_dataset, batch_size=Batch_size, num_workers=0, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, batch_size=Batch_size, num_workers=8,pin_memory=True,
gen_val = DataLoader(val_dataset, batch_size=Batch_size, num_workers=0,pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
else:
gen = Generator(Batch_size, lines[:num_train],
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册