Repository: pyg-team/pytorch_geometric
Pull request #7383 (2023-05-17 20:56:01)

Drop multi-threading in `PrefetchLoader`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7383?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7383](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7383?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (c4da3e4) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/def5301f503bf27c27e8d19bcf0e252275990634?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (def5301) will **increase** coverage by `0.09%`.
> The diff coverage is `80.00%`.

> :exclamation: Current head c4da3e4 differs from pull request most recent head e93ec00. Consider uploading reports for the commit e93ec00 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7383      +/-   ##
==========================================
+ Coverage   91.21%   91.31%   +0.09%     
==========================================
  Files         439      439              
  Lines       24500    24493       -7     
==========================================
+ Hits        22347    22365      +18     
+ Misses       2153     2128      -25     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7383?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/loader/prefetch.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7383?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci9wcmVmZXRjaC5weQ==) | `77.50% <80.00%> (+49.84%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 27c4f2b1c..0ccdb13eb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added an optional `batch_size` argument to `fps`, `knn`, `knn_graph`, `radius` and `radius_graph` ([#7368](https://github.com/pyg-team/pytorch_geometric/pull/7368))
-- Added `GPUPrefetcher` capabilities ([#7376](https://github.com/pyg-team/pytorch_geometric/pull/7376), [#7378](https://github.com/pyg-team/pytorch_geometric/pull/7378))
+- Added `PrefetchLoader` capabilities ([#7376](https://github.com/pyg-team/pytorch_geometric/pull/7376), [#7378](https://github.com/pyg-team/pytorch_geometric/pull/7378), [#7383](https://github.com/pyg-team/pytorch_geometric/pull/7383))
 - Added an example for hierarichial sampling ([#7244](https://github.com/pyg-team/pytorch_geometric/pull/7244))
 - Added Kùzu remote backend examples ([#7298](https://github.com/pyg-team/pytorch_geometric/pull/7298))
 - Fixed tracing of `add_self_loops` for a dynamic number of nodes ([#7330](https://github.com/pyg-team/pytorch_geometric/pull/7330))
diff --git a/torch_geometric/loader/__init__.py b/torch_geometric/loader/__init__.py
index 600d3483e..494a38002 100644
--- a/torch_geometric/loader/__init__.py
+++ b/torch_geometric/loader/__init__.py
@@ -18,6 +18,7 @@ from .temporal_dataloader import TemporalDataLoader
 from .neighbor_sampler import NeighborSampler
 from .imbalanced_sampler import ImbalancedSampler
 from .dynamic_batch_sampler import DynamicBatchSampler
+from .prefetch import PrefetchLoader
 from .mixin import AffinityMixin
 
 __all__ = classes = [
@@ -42,6 +43,7 @@ __all__ = classes = [
     'NeighborSampler',
     'ImbalancedSampler',
     'DynamicBatchSampler',
+    'PrefetchLoader',
     'AffinityMixin',
 ]
 
diff --git a/torch_geometric/loader/prefetch.py b/torch_geometric/loader/prefetch.py
index d6dae8f63..3bbfd69c6 100644
--- a/torch_geometric/loader/prefetch.py
+++ b/torch_geometric/loader/prefetch.py
@@ -1,77 +1,69 @@
-from queue import Queue
-from threading import Thread
+from contextlib import nullcontext
+from functools import partial
 from typing import Any, Optional
 
 import torch
 from torch.utils.data import DataLoader
 
 
-class GPUPrefetcher:
-    r"""A GPU prefetcher class for asynchronously loading data from a
+class PrefetchLoader:
+    r"""A GPU prefetcher class for asynchronously transferring data of a
     :class:`torch.utils.data.DataLoader` from host memory to device memory.
 
     Args:
-        loader (torch.utils.DataLoader): A data loader object.
-        device (torch.device): The CUDA device to load the data to.
-        prefetch_size (int, optional): The number of batches to prefetch at
-            once. (default: :obj:`1`)
+        loader (torch.utils.data.DataLoader): The data loader.
+        device (torch.device, optional): The device to load the data to.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
         loader: DataLoader,
-        device: torch.device,
-        prefetch_size: int = 1,
+        device: Optional[torch.device] = None,
     ):
-        if prefetch_size < 1:
-            raise ValueError(f"'prefetch_size' must be greater than 0 "
-                             f"(got {prefetch_size})")
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
         self.loader = loader
         self.device = torch.device(device)
-        self.prefetch_size = prefetch_size
 
-        self.load_stream = torch.cuda.Stream(device=device)
-        self.queue = Queue(maxsize=prefetch_size)
-        self.worker: Optional[Thread] = None
-
-        self.idx = 0
+        self.is_cuda = torch.cuda.is_available() and self.device.type == 'cuda'
 
     def non_blocking_transfer(self, batch: Any) -> Any:
-        # (Recursive) non-blocking device transfer:
+        if not self.is_cuda:
+            return batch
         if isinstance(batch, (list, tuple)):
             return [self.non_blocking_transfer(v) for v in batch]
         if isinstance(batch, dict):
             return {k: self.non_blocking_transfer(v) for k, v in batch.items()}
 
-        with torch.cuda.stream(self.load_stream):
-            batch = batch.pin_memory()
-            return batch.to(self.device, non_blocking=True)
-
-    def load_loop(self):
-        for batch in self.loader:
-            self.queue.put(self.non_blocking_transfer(batch))
-
-    def __iter__(self) -> 'GPUPrefetcher':
-        is_dead = self.worker is None or not self.worker.is_alive()
-        if is_dead and self.queue.empty() and self.idx == 0:
-            self.worker = Thread(target=self.load_loop)
-            self.worker.daemon = True
-            self.worker.start()
-
-        return self
-
-    def __next__(self) -> Any:
-        is_dead = not self.worker.is_alive()
-        if (is_dead and self.queue.empty()) or self.idx >= len(self):
-            self.idx = 0
-            self.queue.join()
-            self.worker.join()
-            raise StopIteration
-
-        out = self.queue.get()
-        self.queue.task_done()
-        self.idx += 1
-        return out
+        batch = batch.pin_memory()
+        return batch.to(self.device, non_blocking=True)
+
+    def __iter__(self) -> Any:
+        first = True
+        if self.is_cuda:
+            stream = torch.cuda.Stream()
+            stream_context = partial(torch.cuda.stream, stream=stream)
+        else:
+            stream = None
+            stream_context = nullcontext
+
+        for next_batch in self.loader:
+
+            with stream_context():
+                next_batch = self.non_blocking_transfer(next_batch)
+
+            if not first:
+                yield batch  # noqa
+            else:
+                first = False
+
+            if stream is not None:
+                torch.cuda.current_stream().wait_stream(stream)
+
+            batch = next_batch
+
+        yield batch
 
     def __len__(self) -> int:
         return len(self.loader)


Test patch:
diff --git a/test/loader/test_prefetch.py b/test/loader/test_prefetch.py
index 32da97e15..caa7ba3ae 100644
--- a/test/loader/test_prefetch.py
+++ b/test/loader/test_prefetch.py
@@ -1,24 +1,21 @@
 import torch
 
-from torch_geometric.loader import NeighborLoader
-from torch_geometric.loader.prefetch import GPUPrefetcher
+from torch_geometric.loader import NeighborLoader, PrefetchLoader
 from torch_geometric.nn import GraphSAGE
-from torch_geometric.testing import onlyCUDA
+from torch_geometric.testing import withCUDA
 
 
-@onlyCUDA
-def test_gpu_prefetcher():
+@withCUDA
+def test_prefetch_loader(device):
     data = [torch.randn(5, 5) for _ in range(10)]
 
-    loader = GPUPrefetcher(data, device='cuda')
-    assert str(loader).startswith('GPUPrefetcher')
+    loader = PrefetchLoader(data, device=device)
+    assert str(loader).startswith('PrefetchLoader')
     assert len(loader) == 10
 
     for i, batch in enumerate(loader):
-        assert batch.is_cuda
+        assert batch.device == device
         assert torch.equal(batch.cpu(), data[i])
-        assert loader.idx > 0
-    assert loader.idx == 0
 
 
 if __name__ == '__main__':
@@ -28,7 +25,7 @@ if __name__ == '__main__':
     from tqdm import tqdm
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--num_workers', type=int, default=8)
+    parser.add_argument('--num_workers', type=int, default=0)
     args = parser.parse_args()
 
     data = PygNodePropPredDataset('ogbn-products', root='/tmp/ogb')[0]
@@ -46,6 +43,7 @@ if __name__ == '__main__':
         num_neighbors=[10, 10],
         num_workers=args.num_workers,
         filter_per_worker=True,
+        persistent_workers=args.num_workers > 0,
     )
 
     print('Forward pass without prefetching...')
@@ -55,6 +53,6 @@ if __name__ == '__main__':
             model(batch.x, batch.edge_index)
 
     print('Forward pass with prefetching...')
-    for batch in tqdm(GPUPrefetcher(loader, device='cuda')):
+    for batch in tqdm(PrefetchLoader(loader)):
         with torch.no_grad():
             model(batch.x, batch.edge_index)