From 1c0d408fbf74dd9233ee5fe015320df24f9e7ca7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 15 Aug 2019 13:44:42 +0200 Subject: [PATCH] updates --- test.py | 2 +- train.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/test.py b/test.py index f98811fe..a16e57c3 100644 --- a/test.py +++ b/test.py @@ -48,7 +48,7 @@ def test(cfg, dataset = LoadImagesAndLabels(test_path, img_size, batch_size) dataloader = DataLoader(dataset, batch_size=batch_size, - num_workers=os.cpu_count(), + num_workers=min(os.cpu_count(), batch_size), pin_memory=True, collate_fn=dataset.collate_fn) diff --git a/train.py b/train.py index 0d1aa45d..99beaa2e 100644 --- a/train.py +++ b/train.py @@ -172,7 +172,7 @@ def train(cfg, # Dataloader dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, - num_workers=opt.num_workers, + num_workers=min(os.cpu_count(), batch_size), shuffle=not opt.rect, # Shuffle=True unless rectangular training is used pin_memory=True, collate_fn=dataset.collate_fn) @@ -341,7 +341,6 @@ if __name__ == '__main__': parser.add_argument('--rect', action='store_true', help='rectangular training') parser.add_argument('--resume', action='store_true', help='resume training flag') parser.add_argument('--transfer', action='store_true', help='transfer learning flag') - parser.add_argument('--num-workers', type=int, default=os.cpu_count(), help='DataLoader workers') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') parser.add_argument('--notest', action='store_true', help='only test final epoch') parser.add_argument('--xywh', action='store_true', help='use xywh loss instead of GIoU loss')