🐛 [AAAI|Fix] The numworker in dataloader

henrytsui000 · henrytsui000 · commit 8132b7a58950 · 2025-07-14T15:30:10.000+08:00
diff --git a/yolo/aaai.py b/yolo/aaai.py
@@ -207,10 +207,10 @@ def setup(self, stage=None):
         )
 
     def train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=self.train_batch, shuffle=self.cfg.task.data.shuffle)
+        return DataLoader(self.train_dataset, batch_size=self.train_batch, shuffle=self.cfg.task.data.shuffle, num_workers=self.cfg.cpu_num)
 
     def val_dataloader(self):
-        return DataLoader(self.val_dataset, batch_size=self.valid_batch, shuffle=False)
+        return DataLoader(self.val_dataset, batch_size=self.valid_batch, shuffle=False, num_workers=self.cfg.cpu_num)
 
 
 @hydra.main(config_path="config", config_name="aaai_config", version_base=None)

Original file line number	Diff line number	Diff line change
`@@ -207,10 +207,10 @@ def setup(self, stage=None):`
`207`	`207`	`)`
`208`	`208`
`209`	`209`	`def train_dataloader(self):`
`210`		`- return DataLoader(self.train_dataset, batch_size=self.train_batch, shuffle=self.cfg.task.data.shuffle)`
	`210`	`+ return DataLoader(self.train_dataset, batch_size=self.train_batch, shuffle=self.cfg.task.data.shuffle, num_workers=self.cfg.cpu_num)`
`211`	`211`
`212`	`212`	`def val_dataloader(self):`
`213`		`- return DataLoader(self.val_dataset, batch_size=self.valid_batch, shuffle=False)`
	`213`	`+ return DataLoader(self.val_dataset, batch_size=self.valid_batch, shuffle=False, num_workers=self.cfg.cpu_num)`
`214`	`214`
`215`	`215`
`216`	`216`	`@hydra.main(config_path="config", config_name="aaai_config", version_base=None)`