Enable multi-process DataLoader for dreambooth (huggingface#950)

skirsten · web-flow · commit 8332c1a6d956 · 2022-10-26T17:24:48.000+02:00
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
@@ -502,7 +502,7 @@ def collate_fn(examples):
         return batch
 
     train_dataloader = torch.utils.data.DataLoader(
-        train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn
+        train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn, num_workers=1
     )
 
     # Scheduler and math around the number of training steps.

Original file line number	Diff line number	Diff line change
`@@ -502,7 +502,7 @@ def collate_fn(examples):`
`502`	`502`	`return batch`
`503`	`503`
`504`	`504`	`train_dataloader = torch.utils.data.DataLoader(`
`505`		`- train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn`
	`505`	`+ train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn, num_workers=1`
`506`	`506`	`)`
`507`	`507`
`508`	`508`	`# Scheduler and math around the number of training steps.`