This commit is contained in:
Glenn Jocher 2019-08-15 13:44:42 +02:00
parent 48af6d136f
commit 1c0d408fbf
2 changed files with 2 additions and 3 deletions

View File

@ -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)

View File

@ -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')