prefetch mvp

catalyst/data/__init__.py

@@ -10,7 +10,10 @@
     MetricLearningTrainDataset,
     QueryGalleryDataset,
 )
-from catalyst.data.loader import BatchLimitLoaderWrapper
+from catalyst.data.loader import (
+    BatchLimitLoaderWrapper,
+    BatchPrefetchLoaderWrapper,
+)
 from catalyst.data.reader import (
     ReaderSpec,
     ScalarReader,

catalyst/data/loader.py

@@ -1,9 +1,53 @@
-from typing import Union
+from typing import Any, Callable, Iterable, Union
+import queue
+import sys
+import threading
 
+import numpy as np
+
+import torch
 from torch.utils.data import DataLoader
 
 
-class BatchLimitLoaderWrapper:
+class ILoaderWrapper:
+    def __init__(self, loader: DataLoader):
+        self.origin = loader
+
+    def __getattr__(self, key):
+        """
+        Gets attribute by ``key``.
+        Firstly, looks at the ``origin`` for the appropriate ``key``.
+        If none founds - looks at the wrappers attributes.
+        If could not found anything - raises ``NotImplementedError``.
+
+        Args:
+            key: attribute key
+
+        Returns:
+            attribute value
+
+        Raises:
+            NotImplementedError: if could not find attribute in ``origin``
+                or ``wrapper``
+        """
+        value = getattr(self.origin, key, None)
+        if value is not None:
+            return value
+        value = getattr(self, key, None)
+        if value is not None:
+            return value
+        raise NotImplementedError()
+
+    def __len__(self) -> int:
+        """Returns length of the wrapper loader.
+
+        Returns:
+            int: length of the wrapper loader
+        """
+        return len(self.origin)
+
+
+class BatchLimitLoaderWrapper(ILoaderWrapper):
     """
     Loader wrapper. Limits number of batches used per each iteration.
 
@@ -50,6 +94,7 @@ def __init__(self, loader: DataLoader, num_batches: Union[int, float]):
             num_batches (Union[int, float]): number of batches to use (int),
                 or portion of iterator (float, should be in [0;1] range)
         """
+        super().__init__(loader)
         assert isinstance(num_batches, (int, float)), (
             "Expected ``num_batches`` type is int/float"
             f"but got {type(num_batches)}"
@@ -61,36 +106,10 @@ def __init__(self, loader: DataLoader, num_batches: Union[int, float]):
             )
             num_batches = int(len(loader) * num_batches)
 
-        self.origin = loader
         self.iterator = iter(self.origin)
         self.iteration_index = 0
         self.num_batches = num_batches
 
-    def __getattr__(self, key):
-        """
-        Gets attribute by ``key``.
-        Firstly, looks at the ``origin`` for the appropriate ``key``.
-        If none founds - looks at the wrappers attributes.
-        If could not found anything - raises ``NotImplementedError``.
-
-        Args:
-            key: attribute key
-
-        Returns:
-            attribute value
-
-        Raises:
-            NotImplementedError: if could not find attribute in ``origin``
-                or ``wrapper``
-        """
-        value = getattr(self.origin, key, None)
-        if value is not None:
-            return value
-        value = getattr(self, key, None)
-        if value is not None:
-            return value
-        raise NotImplementedError()
-
     def __iter__(self):
         """Iterator.
 
@@ -115,13 +134,85 @@ def __next__(self):
         batch = next(self.iterator)
         return batch
 
-    def __len__(self) -> int:
-        """Returns length of the wrapper loader.
 
-        Returns:
-            int: length of the wrapper loader
-        """
-        return len(self.origin)
+def _any2cuda_non_blocking(value: Any):
+    # based on catalyst.utils.torch.any2device
+    # but with cuda non_blocking trick
+    if isinstance(value, dict):
+        return {k: _any2cuda_non_blocking(v) for k, v in value.items()}
+    elif isinstance(value, (tuple, list)):
+        return [_any2cuda_non_blocking(v) for v in value]
+    elif torch.is_tensor(value):
+        return value.cuda(non_blocking=True)
+    elif (
+        isinstance(value, (np.ndarray, np.void))
+        and value.dtype.fields is not None
+    ):
+        return {
+            k: _any2cuda_non_blocking(value[k])
+            for k in value.dtype.fields.keys()
+        }
+    elif isinstance(value, np.ndarray):
+        return torch.tensor(value).cuda(non_blocking=True)
+
+
+def _map_loop(
+    func: Callable,
+    iterable: Iterable,
+    result_queue: queue.Queue,
+    error_queue: queue.Queue,
+    done_event: threading.Event,
+):
+    try:
+        for x in iterable:
+            result = func(x)
+            result_queue.put(result)
+    except BaseException:
+        error_queue.put(sys.exc_info())
+    finally:
+        done_event.set()
+
+
+def _prefetch_map(
+    func: Callable,
+    iterable: Iterable,
+    num_prefetches: int = 1,
+    timeout: int = 2,
+) -> Iterable:
+    result_queue = queue.Queue(num_prefetches)
+    error_queue = queue.Queue(1)
+    done_event = threading.Event()
+    map_thread = threading.Thread(
+        target=_map_loop,
+        args=(func, iterable, result_queue, error_queue, done_event),
+    )
+    map_thread.daemon = True
+    map_thread.start()
+    while not (done_event.is_set() and result_queue.empty()):
+        try:
+            result = result_queue.get(timeout=timeout)
+        except queue.Empty:
+            continue
+        yield result
+    if error_queue.full():
+        raise error_queue.get()[1]
+
+
+def _prefetch_loader(loader: DataLoader, num_prefetches: int) -> Iterable:
+    if torch.cuda.is_available():
+        loader = _prefetch_map(
+            _any2cuda_non_blocking, loader, num_prefetches=num_prefetches,
+        )
+    return loader
+
+
+class BatchPrefetchLoaderWrapper(ILoaderWrapper):
+    def __init__(self, loader: DataLoader, num_prefetches: int = None):
+        super().__init__(loader)
+        self.num_prefetches = num_prefetches or loader.batch_size
+
+    def __iter__(self):
+        return _prefetch_loader(self.origin, self.num_prefetches)
 
 
-__all__ = ["BatchLimitLoaderWrapper"]
+__all__ = ["BatchLimitLoaderWrapper", "BatchPrefetchLoaderWrapper"]

catalyst/utils/torch.py

@@ -143,7 +143,7 @@ def any2device(value, device: Device):
             k: any2device(value[k], device) for k in value.dtype.fields.keys()
         }
     elif isinstance(value, np.ndarray):
-        return torch.Tensor(value).to(device)
+        return torch.tensor(value, device=device)
     return value