Add --prefetch-factor option in trainer

KCBF · Feb 7, 2023 · 42eecf6 · 42eecf6
1 parent fd8f8d1
commit 42eecf6
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/imagenet/training/trainer.py b/imagenet/training/trainer.py
@@ -32,6 +32,7 @@ def create_dataloader(self, type):
                 shuffle=False,
                 pin_memory=True,
                 num_workers=self.args.num_workers,
+                prefetch_factor=self.args.prefetch_factor,
                 drop_last=True)
             return loader
         else:
@@ -45,6 +46,7 @@ def create_dataloader(self, type):
                 shuffle=False,
                 pin_memory=True,
                 num_workers=self.args.num_workers,
+                prefetch_factor=self.args.prefetch_factor,
                 drop_last=False)
             return loader
 

diff --git a/nunif/training/trainer.py b/nunif/training/trainer.py
@@ -201,6 +201,8 @@ def create_trainer_default_parser():
                         help="momentum for sgd")
     parser.add_argument("--num-workers", type=int, default=num_workers,
                         help="number of worker processes for data loader")
+    parser.add_argument("--prefetch-factor", type=int, default=4,
+                        help="number of batches loaded in advance by each worker")
     parser.add_argument("--max-epoch", type=int, default=200,
                         help="max epoch")
     parser.add_argument("--gpu", type=int, nargs="+", default=[0],

diff --git a/waifu2x/training/trainer.py b/waifu2x/training/trainer.py
@@ -78,6 +78,7 @@ def create_dataloader(self, type):
                 pin_memory=True,
                 sampler=dataset.sampler(),
                 num_workers=self.args.num_workers,
+                prefetch_factor=self.args.prefetch_factor,
                 drop_last=True)
         elif type == "eval":
             dataset = Waifu2xDataset(
@@ -94,6 +95,7 @@ def create_dataloader(self, type):
                 worker_init_fn=dataset.worker_init,
                 shuffle=False,
                 num_workers=self.args.num_workers,
+                prefetch_factor=self.args.prefetch_factor,
                 drop_last=False)
 
     def create_env(self):