This commit is contained in:
Glenn Jocher 2019-09-19 02:10:55 +02:00
parent 2f436d499a
commit 6d8e82c175
2 changed files with 2 additions and 2 deletions

View File

@ -48,7 +48,7 @@ def test(cfg,
dataset = LoadImagesAndLabels(test_path, img_size, batch_size) dataset = LoadImagesAndLabels(test_path, img_size, batch_size)
dataloader = DataLoader(dataset, dataloader = DataLoader(dataset,
batch_size=batch_size, batch_size=batch_size,
num_workers=min(os.cpu_count(), batch_size), num_workers=min([os.cpu_count(), batch_size, 16]),
pin_memory=True, pin_memory=True,
collate_fn=dataset.collate_fn) collate_fn=dataset.collate_fn)

View File

@ -194,7 +194,7 @@ def train():
# Dataloader # Dataloader
dataloader = torch.utils.data.DataLoader(dataset, dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size, batch_size=batch_size,
num_workers=min(os.cpu_count(), batch_size), num_workers=min([os.cpu_count(), batch_size, 16]),
shuffle=not opt.rect, # Shuffle=True unless rectangular training is used shuffle=not opt.rect, # Shuffle=True unless rectangular training is used
pin_memory=True, pin_memory=True,
collate_fn=dataset.collate_fn) collate_fn=dataset.collate_fn)