updates
This commit is contained in:
parent
54daa69adb
commit
a2dc8a6b5a
4
train.py
4
train.py
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue