Repository: pyg-team/pytorch_geometric
Pull request #6829 (2023-03-01 16:06:33)

Added `ZipLoader` to execute multiple `NodeLoader` or `LinkLoader` instances

codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6829](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (06b8df7) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/44875887afab882a3a44bdc11dfe3fe4c2cc3104?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4487588) will **decrease** coverage by `0.32%`.
> The diff coverage is `17.94%`.

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

```diff
@@            Coverage Diff             @@
##           master    #6829      +/-   ##
==========================================
- Coverage   91.65%   91.33%   -0.32%     
==========================================
  Files         426      427       +1     
  Lines       23050    23047       -3     
==========================================
- Hits        21126    21050      -76     
- Misses       1924     1997      +73     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?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/zip\_loader.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci96aXBfbG9hZGVyLnB5) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/loader/link\_loader.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci9saW5rX2xvYWRlci5weQ==) | `93.68% <100.00%> (+0.06%)` | :arrow_up: |
| [torch\_geometric/loader/node\_loader.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci9ub2RlX2xvYWRlci5weQ==) | `91.54% <100.00%> (+0.12%)` | :arrow_up: |
| [torch\_geometric/sampler/neighbor\_sampler.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvbmVpZ2hib3Jfc2FtcGxlci5weQ==) | `91.56% <100.00%> (-3.22%)` | :arrow_down: |
| [torch\_geometric/nn/conv/utils/typing.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvdXRpbHMvdHlwaW5nLnB5) | `83.75% <0.00%> (-15.00%)` | :arrow_down: |
| [torch\_geometric/nn/pool/asap.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvYXNhcC5weQ==) | `91.89% <0.00%> (-8.11%)` | :arrow_down: |
| [torch\_geometric/nn/dense/linear.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2RlbnNlL2xpbmVhci5weQ==) | `88.97% <0.00%> (-5.92%)` | :arrow_down: |
| [torch\_geometric/nn/models/attentive\_fp.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9hdHRlbnRpdmVfZnAucHk=) | `95.83% <0.00%> (-4.17%)` | :arrow_down: |
| [torch\_geometric/io/tu.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2lvL3R1LnB5) | `90.80% <0.00%> (-2.30%)` | :arrow_down: |
| ... and [28 more](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6829?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | |

: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 288b6d4b4..bad092056 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,8 +7,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `ZipLoader` to execute multiple `NodeLoader` or `LinkLoader` instances ([#6829](https://github.com/pyg-team/pytorch_geometric/issues/6829))
 - Added common `utils.select` and `utils.narrow` functionality to support filtering of both tensors and lists ([#6162](https://github.com/pyg-team/pytorch_geometric/issues/6162))
-- Support `normalization` customization in `get_mesh_laplacian` ([#6037](https://github.com/pyg-team/pytorch_geometric/issues/6037))
+- Support `normalization` customization in `get_mesh_laplacian` ([#6790](https://github.com/pyg-team/pytorch_geometric/issues/6790))
 - Added the `TemporalEncoding` module ([#6785](https://github.com/pyg-team/pytorch_geometric/pull/6785))
 - Added CPU-optimized `spmm_reduce` functionality via CSR format ([#6699](https://github.com/pyg-team/pytorch_geometric/pull/6699), [#6759](https://github.com/pyg-team/pytorch_geometric/pull/6759))
 - Added support for the revised version of the `MD17` dataset ([#6734](https://github.com/pyg-team/pytorch_geometric/pull/6734))
@@ -93,7 +94,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
-- Fixed a bug in `get_mesh_laplacian` ([#6037](https://github.com/pyg-team/pytorch_geometric/issues/6037))
+- Fixed a bug in `get_mesh_laplacian` ([#6790](https://github.com/pyg-team/pytorch_geometric/issues/6790))
 - Fixed a bug in which masks were not properly masked in `GNNExplainer` on link prediction tasks ([#6787](https://github.com/pyg-team/pytorch_geometric/pull/6787))
 - Allow the usage of `ChebConv` within `GNNExplainer` ([#6778](https://github.com/pyg-team/pytorch_geometric/pull/6778))
 - Allow setting the `EdgeStorage.num_edges` property ([#6710](https://github.com/pyg-team/pytorch_geometric/pull/6710))
diff --git a/torch_geometric/loader/__init__.py b/torch_geometric/loader/__init__.py
index 1e2d7737b..b66031413 100644
--- a/torch_geometric/loader/__init__.py
+++ b/torch_geometric/loader/__init__.py
@@ -11,6 +11,7 @@ from .graph_saint import (GraphSAINTSampler, GraphSAINTNodeSampler,
                           GraphSAINTEdgeSampler, GraphSAINTRandomWalkSampler)
 from .shadow import ShaDowKHopSampler
 from .random_node_loader import RandomNodeLoader
+from .zip_loader import ZipLoader
 from .data_list_loader import DataListLoader
 from .dense_data_loader import DenseDataLoader
 from .temporal_dataloader import TemporalDataLoader
@@ -33,6 +34,7 @@ __all__ = classes = [
     'GraphSAINTRandomWalkSampler',
     'ShaDowKHopSampler',
     'RandomNodeLoader',
+    'ZipLoader',
     'DataListLoader',
     'DenseDataLoader',
     'TemporalDataLoader',
diff --git a/torch_geometric/loader/link_loader.py b/torch_geometric/loader/link_loader.py
index b619d56e4..3f0476769 100644
--- a/torch_geometric/loader/link_loader.py
+++ b/torch_geometric/loader/link_loader.py
@@ -1,6 +1,7 @@
 from typing import Any, Callable, Iterator, List, Optional, Tuple, Union
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
 from torch_geometric.loader.base import DataLoaderIterator
@@ -176,7 +177,7 @@ class LinkLoader(torch.utils.data.DataLoader, AffinityMixin):
         iterator = range(edge_label_index.size(1))
         super().__init__(iterator, collate_fn=self.collate_fn, **kwargs)
 
-    def collate_fn(self, index: List[int]) -> Any:
+    def collate_fn(self, index: Union[Tensor, List[int]]) -> Any:
         r"""Samples a subgraph from a batch of input nodes."""
         input_data: EdgeSamplerInput = self.input_data[index]
 
diff --git a/torch_geometric/loader/node_loader.py b/torch_geometric/loader/node_loader.py
index d1713c738..b64687f57 100644
--- a/torch_geometric/loader/node_loader.py
+++ b/torch_geometric/loader/node_loader.py
@@ -1,6 +1,7 @@
-from typing import Any, Callable, Iterator, Optional, Tuple, Union
+from typing import Any, Callable, Iterator, List, Optional, Tuple, Union
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
 from torch_geometric.loader.base import DataLoaderIterator
@@ -111,7 +112,7 @@ class NodeLoader(torch.utils.data.DataLoader, AffinityMixin):
         iterator = range(input_nodes.size(0))
         super().__init__(iterator, collate_fn=self.collate_fn, **kwargs)
 
-    def collate_fn(self, index: NodeSamplerInput) -> Any:
+    def collate_fn(self, index: Union[Tensor, List[int]]) -> Any:
         r"""Samples a subgraph from a batch of input nodes."""
         input_data: NodeSamplerInput = self.input_data[index]
 
diff --git a/torch_geometric/loader/zip_loader.py b/torch_geometric/loader/zip_loader.py
new file mode 100644
index 000000000..cb76c9d11
--- /dev/null
+++ b/torch_geometric/loader/zip_loader.py
@@ -0,0 +1,78 @@
+from typing import Any, Iterator, List, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.loader import LinkLoader, NodeLoader
+from torch_geometric.loader.base import DataLoaderIterator
+
+
+class ZipLoader(torch.utils.data.DataLoader):
+    r"""A loader that returns a tuple of data objects by sampling from multiple
+    :class:`NodeLoader` or :class:`LinkLoader` instances.
+
+    Args:
+        loaders (List[NodeLoader] or List[LinkLoader]): The loader instances.
+        filter_per_worker (bool, optional): If set to :obj:`True`, will filter
+            the returning data in each worker's subprocess rather than in the
+            main process.
+            Setting this to :obj:`True` for in-memory datasets is generally not
+            recommended:
+            (1) it may result in too many open file handles,
+            (2) it may slown down data loading,
+            (3) it requires operating on CPU tensors.
+            (default: :obj:`False`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
+            :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.
+    """
+    def __init__(
+        self,
+        loaders: Union[List[NodeLoader], List[LinkLoader]],
+        filter_per_worker: bool = False,
+        **kwargs,
+    ):
+        # Remove for PyTorch Lightning:
+        if 'dataset' in kwargs:
+            del kwargs['dataset']
+        if 'collate_fn' in kwargs:
+            del kwargs['collate_fn']
+
+        for loader in loaders:
+            if not callable(getattr(loader, 'collate_fn', None)):
+                raise ValueError("'{loader.__class__.__name__}' does not have "
+                                 "a 'collate_fn' method")
+            if not callable(getattr(loader, 'filter_fn', None)):
+                raise ValueError("'{loader.__class__.__name__}' does not have "
+                                 "a 'filter_fn' method")
+            loader.filter_per_worker = filter_per_worker
+
+        iterator = range(min([len(loader.dataset) for loader in loaders]))
+        super().__init__(iterator, collate_fn=self.collate_fn, **kwargs)
+
+        self.loaders = loaders
+        self.filter_per_worker = filter_per_worker
+
+    def collate_fn(self, index: List[int]) -> Tuple[Any, ...]:
+        if not isinstance(index, Tensor):
+            index = torch.tensor(index, dtype=torch.long)
+
+        return tuple(loader.collate_fn(index) for loader in self.loaders)
+
+    def filter_fn(
+        self,
+        outs: Tuple[Any, ...],
+    ) -> Tuple[Union[Data, HeteroData], ...]:
+        loaders = self.loaders
+        return tuple(loader.filter_fn(v) for loader, v in zip(loaders, outs))
+
+    def _get_iterator(self) -> Iterator:
+        if self.filter_per_worker:
+            return super()._get_iterator()
+
+        # Execute `filter_fn` in the main process:
+        return DataLoaderIterator(super()._get_iterator(), self.filter_fn)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(loaders={self.loaders})'


Test patch:
diff --git a/test/loader/test_zip_loader.py b/test/loader/test_zip_loader.py
new file mode 100644
index 000000000..14c4c36d4
--- /dev/null
+++ b/test/loader/test_zip_loader.py
@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.loader import NeighborLoader, ZipLoader
+
+
+@pytest.mark.parametrize('filter_per_worker', [True, False])
+def test_zip_loader(filter_per_worker):
+    x = torch.arange(100)
+    edge_index = torch.randint(0, 100, (2, 1000))
+    data = Data(x=x, edge_index=edge_index)
+
+    loaders = [
+        NeighborLoader(data, [5], input_nodes=torch.arange(0, 50)),
+        NeighborLoader(data, [5], input_nodes=torch.arange(50, 95)),
+    ]
+
+    loader = ZipLoader(loaders, batch_size=10,
+                       filter_per_worker=filter_per_worker)
+
+    assert str(loader) == ('ZipLoader(loaders=[NeighborLoader(), '
+                           'NeighborLoader()])')
+    assert len(loader) == 5
+    assert loader.dataset == range(0, 45)
+
+    for i, (batch1, batch2) in enumerate(loader):
+        n_id1 = batch1.n_id[:batch1.batch_size]
+        n_id2 = batch2.n_id[:batch2.batch_size]
+
+        if i < 4:
+            assert batch1.batch_size == 10
+            assert batch2.batch_size == 10
+            assert torch.equal(n_id1, torch.arange(0 + i * 10, 10 + i * 10))
+            assert torch.equal(n_id2, torch.arange(50 + i * 10, 60 + i * 10))
+        else:
+            assert batch1.batch_size == 5
+            assert batch2.batch_size == 5
+            assert torch.equal(n_id1, torch.arange(0 + i * 10, 5 + i * 10))
+            assert torch.equal(n_id2, torch.arange(50 + i * 10, 55 + i * 10))
diff --git a/test/nn/test_inits.py b/test/nn/test_inits.py
index 113b8cbb6..6d402f8f5 100644
--- a/test/nn/test_inits.py
+++ b/test/nn/test_inits.py
@@ -21,12 +21,12 @@ def test_inits():
     assert x.max() <= 0.5
 
     glorot(x)
-    assert x.min() >= -1.25
-    assert x.max() <= 1.25
+    assert x.min() >= -2.0
+    assert x.max() <= 2.0
 
     glorot_orthogonal(x, scale=1.0)
-    assert x.min() >= -1.25
-    assert x.max() <= 1.25
+    assert x.min() >= -2.0
+    assert x.max() <= 2.0
 
     zeros(x)
     assert x.tolist() == [[0, 0, 0, 0]]