2X training speed improvement: prefetch next batch from the data loader while running current batch

src/super_gradients/training/sg_trainer/sg_trainer.py

@@ -5,6 +5,7 @@
 import warnings
 from copy import deepcopy
 from typing import Union, Tuple, Mapping, Dict, Any, List, Optional
+import concurrent.futures
 
 import hydra
 import numpy as np
@@ -116,6 +117,32 @@
 logger = get_logger(__name__)
 
 
+class PrefetchIterable:
+    def __init__(self, iterable):
+        self.iterable = iterable
+
+    def __iter__(self):
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+
+        try:
+            iterator = iter(self.iterable)
+
+            def _prefetch():
+                return next(iterator)
+
+            prefetch_future = executor.submit(_prefetch)
+            while True:
+                try:
+                    value = prefetch_future.result()
+                except StopIteration:
+                    return
+                prefetch_future = executor.submit(_prefetch)
+                yield value
+
+        finally:
+            executor.shutdown()
+
+
 class Trainer:
     """
     SuperGradient Model - Base Class for Sg Models
@@ -471,7 +498,7 @@ def _train_epoch(self, context: PhaseContext, silent_mode: bool = False) -> tupl
 
         # THE DISABLE FLAG CONTROLS WHETHER THE PROGRESS BAR IS SILENT OR PRINTS THE LOGS
         with tqdm(
-            self.train_loader, total=expected_iterations, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True, disable=silent_mode
+            PrefetchIterable(self.train_loader), total=expected_iterations, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True, disable=silent_mode
         ) as progress_bar_train_loader:
             progress_bar_train_loader.set_description(f"Train epoch {context.epoch}")
 
@@ -2258,7 +2285,7 @@ def evaluate(
         expected_iterations = len(data_loader) if max_batches is None else max_batches
 
         with tqdm(
-            data_loader, total=expected_iterations, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True, disable=silent_mode
+            PrefetchIterable(data_loader), total=expected_iterations, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True, disable=silent_mode
         ) as progress_bar_data_loader:
             if not silent_mode:
                 # PRINT TITLES