This commit is contained in:
Glenn Jocher 2019-12-04 15:15:23 -08:00
parent 54daa69adb
commit a2dc8a6b5a
1 changed files with 3 additions and 1 deletions

View File

@ -199,9 +199,11 @@ def train():
# Dataloader # Dataloader
batch_size = min(batch_size, len(dataset)) batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 16])
print('Using %g dataloader workers' % nw)
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 if batch_size > 1 else 0, 16]), num_workers=nw,
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)