From a7f570c9b72369fd75e15733d5645d09039d5f9e Mon Sep 17 00:00:00 2001
From: "AFLALO, Jonathan Isaac" <jonathan.aflalo@alibaba-inc.com>
Date: Tue, 5 May 2020 14:32:10 +0300
Subject: [PATCH] added MultiEpochsDataLoader

---
 timm/data/loader.py | 40 +++++++++++++++++++++++++++++++++++++++-
 train.py            |  3 +++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/timm/data/loader.py b/timm/data/loader.py
index f3faf7b9..4dc0d697 100644
--- a/timm/data/loader.py
+++ b/timm/data/loader.py
@@ -140,6 +140,7 @@ def create_loader(
         pin_memory=False,
         fp16=False,
         tf_preprocessing=False,
+        use_multi_epochs_loader=False
 ):
     re_num_splits = 0
     if re_split:
@@ -175,7 +176,12 @@ def create_loader(
     if collate_fn is None:
         collate_fn = fast_collate if use_prefetcher else torch.utils.data.dataloader.default_collate
 
-    loader = torch.utils.data.DataLoader(
+    loader_class = torch.utils.data.DataLoader
+
+    if use_multi_epochs_loader:
+        loader_class = MultiEpochsDataLoader
+
+    loader = loader_class(
         dataset,
         batch_size=batch_size,
         shuffle=sampler is None and is_training,
@@ -198,3 +204,35 @@ def create_loader(
         )
 
     return loader
+
+
+class MultiEpochsDataLoader(torch.utils.data.DataLoader):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._DataLoader__initialized = False
+        self.batch_sampler = _RepeatSampler(self.batch_sampler)
+        self._DataLoader__initialized = True
+        self.iterator = super().__iter__()
+
+    def __len__(self):
+        return len(self.batch_sampler.sampler)
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield next(self.iterator)
+
+
+class _RepeatSampler(object):
+    """ Sampler that repeats forever.
+
+    Args:
+        sampler (Sampler)
+    """
+
+    def __init__(self, sampler):
+        self.sampler = sampler
+
+    def __iter__(self):
+        while True:
+            yield from iter(self.sampler)
diff --git a/train.py b/train.py
index e88640a7..899c6984 100755
--- a/train.py
+++ b/train.py
@@ -198,6 +198,8 @@ parser.add_argument('--eval-metric', default='top1', type=str, metavar='EVAL_MET
 parser.add_argument('--tta', type=int, default=0, metavar='N',
                     help='Test/inference time augmentation (oversampling) factor. 0=None (default: 0)')
 parser.add_argument("--local_rank", default=0, type=int)
+parser.add_argument('--use-multi-epochs-loader', action='store_true', default=False,
+                    help='use the multi-epochs-loader to save time at the beginning of every epoch')
 
 
 def _parse_args():
@@ -391,6 +393,7 @@ def main():
         distributed=args.distributed,
         collate_fn=collate_fn,
         pin_memory=args.pin_mem,
+        use_multi_epochs_loader=args.use_multi_epochs_loader
     )
 
     eval_dir = os.path.join(args.data, 'val')