Support np.memmap in NeighborLoader (#5696)

Author: rusty1s

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.2.0] - 2022-MM-DD
 ### Added
+- Added `np.memmap` support in `NeighborLoader` ([#5696](https://github.com/pyg-team/pytorch_geometric/pull/5696))
 - Added `assortativity` that computes degree assortativity coefficient ([#5587](https://github.com/pyg-team/pytorch_geometric/pull/5587))
 - Added `SSGConv` layer ([#5599](https://github.com/pyg-team/pytorch_geometric/pull/5599))
 - Added `shuffle_node`, `mask_feature` and `add_random_edge` augmentation methdos ([#5548](https://github.com/pyg-team/pytorch_geometric/pull/5548))

test/data/test_batch.py

@@ -196,7 +196,7 @@ def test_pickling():
     assert id(batch._store._parent()) == id(batch)
     assert batch.num_nodes == 20
 
-    path = f'{random.randrange(sys.maxsize)}.pt'
+    path = os.path.join(os.sep, 'tmp', f'{random.randrange(sys.maxsize)}.pt')
     torch.save(batch, path)
     assert id(batch._store._parent()) == id(batch)
     assert batch.num_nodes == 20

test/loader/test_neighbor_loader.py

@@ -1,3 +1,7 @@
+import os
+import random
+import sys
+
 import numpy as np
 import pytest
 import torch
@@ -516,3 +520,25 @@ def test_pyg_lib_heterogeneous_neighbor_loader():
     assert len(edge_id1_dict) == len(edge_id2_dict)
     for key in edge_id1_dict.keys():
         assert torch.equal(edge_id1_dict[key], edge_id2_dict[key])
+
+
+def test_memmap_neighbor_loader():
+    path = os.path.join(os.sep, 'tmp', f'{random.randrange(sys.maxsize)}.npy')
+    x = np.memmap(path, dtype=np.float32, mode='w+', shape=(100, 32))
+    x[:] = np.random.randn(100, 32)
+
+    data = Data()
+    data.x = np.memmap(path, dtype=np.float32, mode='r', shape=(100, 32))
+    data.edge_index = get_edge_index(100, 100, 500)
+
+    assert str(data) == 'Data(x=[100, 32], edge_index=[2, 500])'
+    assert data.num_nodes == 100
+
+    loader = NeighborLoader(data, num_neighbors=[5] * 2, batch_size=20,
+                            num_workers=6)
+    batch = next(iter(loader))
+    assert batch.num_nodes <= 100
+    assert isinstance(batch.x, torch.Tensor)
+    assert batch.x.size() == (batch.num_nodes, 32)
+
+    os.remove(path)

torch_geometric/data/storage.py

@@ -15,6 +15,7 @@
     Union,
 )
 
+import numpy as np
 import torch
 from torch import Tensor
 from torch_sparse import SparseTensor, coalesce
@@ -265,8 +266,10 @@ def num_nodes(self) -> Optional[int]:
         if 'num_nodes' in self:
             return self['num_nodes']
         for key, value in self.items():
-            if isinstance(value, Tensor) and (key in N_KEYS or 'node' in key):
-                return value.size(self._parent().__cat_dim__(key, value, self))
+            if (isinstance(value, (Tensor, np.ndarray))
+                    and (key in N_KEYS or 'node' in key)):
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.shape[cat_dim]
         if 'adj' in self and isinstance(self.adj, SparseTensor):
             return self.adj.size(0)
         if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):
@@ -291,7 +294,9 @@ def num_nodes(self) -> Optional[int]:
 
     @property
     def num_node_features(self) -> int:
-        if 'x' in self and isinstance(self.x, (Tensor, SparseTensor)):
+        if 'x' in self and isinstance(self.x, (Tensor, np.ndarray)):
+            return 1 if self.x.ndim == 1 else self.x.shape[-1]
+        if 'x' in self and isinstance(self.x, SparseTensor):
             return 1 if self.x.dim() == 1 else self.x.size(-1)
         return 0
 
@@ -302,9 +307,9 @@ def num_features(self) -> int:
     def is_node_attr(self, key: str) -> bool:
         value = self[key]
         cat_dim = self._parent().__cat_dim__(key, value, self)
-        if not isinstance(value, Tensor):
+        if not isinstance(value, (Tensor, np.ndarray)):
             return False
-        if value.dim() == 0 or value.size(cat_dim) != self.num_nodes:
+        if value.ndim == 0 or value.shape[cat_dim] != self.num_nodes:
             return False
         return True
 
@@ -350,17 +355,19 @@ def edge_index(self) -> Tensor:
     def num_edges(self) -> int:
         # We sequentially access attributes that reveal the number of edges.
         for key, value in self.items():
-            if isinstance(value, Tensor) and 'edge' in key:
-                return value.size(self._parent().__cat_dim__(key, value, self))
+            if isinstance(value, (Tensor, np.ndarray)) and 'edge' in key:
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.shape[cat_dim]
         for value in self.values('adj', 'adj_t'):
             if isinstance(value, SparseTensor):
                 return value.nnz()
         return 0
 
     @property
     def num_edge_features(self) -> int:
-        if 'edge_attr' in self and isinstance(self.edge_attr, Tensor):
-            return 1 if self.edge_attr.dim() == 1 else self.edge_attr.size(-1)
+        if ('edge_attr' in self and isinstance(self.edge_attr,
+                                               (Tensor, np.ndarray))):
+            return 1 if self.edge_attr.ndim == 1 else self.edge_attr.shape[-1]
         return 0
 
     @property
@@ -386,9 +393,9 @@ def is_node_attr(self, key: str) -> bool:
     def is_edge_attr(self, key: str) -> bool:
         value = self[key]
         cat_dim = self._parent().__cat_dim__(key, value, self)
-        if not isinstance(value, Tensor):
+        if not isinstance(value, (Tensor, np.ndarray)):
             return False
-        if value.dim() == 0 or value.size(cat_dim) != self.num_edges:
+        if value.ndim == 0 or value.shape[cat_dim] != self.num_edges:
             return False
         return True
 
@@ -467,9 +474,9 @@ def is_node_attr(self, key: str) -> bool:
         cat_dim = self._parent().__cat_dim__(key, value, self)
 
         num_nodes, num_edges = self.num_nodes, self.num_edges
-        if not isinstance(value, Tensor):
+        if not isinstance(value, (Tensor, np.ndarray)):
             return False
-        if value.dim() == 0 or value.size(cat_dim) != num_nodes:
+        if value.ndim == 0 or value.shape[cat_dim] != num_nodes:
             return False
         if num_nodes != num_edges:
             return True
@@ -480,9 +487,9 @@ def is_edge_attr(self, key: str) -> bool:
         cat_dim = self._parent().__cat_dim__(key, value, self)
 
         num_nodes, num_edges = self.num_nodes, self.num_edges
-        if not isinstance(value, Tensor):
+        if not isinstance(value, (Tensor, np.ndarray)):
             return False
-        if value.dim() == 0 or value.size(cat_dim) != num_edges:
+        if value.ndim == 0 or value.shape[cat_dim] != num_edges:
             return False
         if num_nodes != num_edges:
             return True

torch_geometric/loader/utils.py

@@ -3,6 +3,7 @@
 from collections.abc import Sequence
 from typing import Dict, Optional, Tuple, Union
 
+import numpy as np
 import torch
 from torch import Tensor
 from torch_sparse import SparseTensor
@@ -11,20 +12,34 @@
 from torch_geometric.data.feature_store import FeatureStore, TensorAttr
 from torch_geometric.data.graph_store import GraphStore
 from torch_geometric.data.storage import EdgeStorage, NodeStorage
-from torch_geometric.typing import InputEdges, InputNodes, OptTensor
+from torch_geometric.typing import (
+    FeatureTensorType,
+    InputEdges,
+    InputNodes,
+    OptTensor,
+)
 
 
-def index_select(value: Tensor, index: Tensor, dim: int = 0) -> Tensor:
-    out: Optional[Tensor] = None
-    if torch.utils.data.get_worker_info() is not None:
-        # If we are in a background process, we write directly into a shared
-        # memory tensor to avoid an extra copy:
-        size = list(value.size())
-        size[dim] = index.numel()
-        numel = math.prod(size)
-        storage = value.storage()._new_shared(numel)
-        out = value.new(storage).view(size)
-    return torch.index_select(value, dim, index, out=out)
+def index_select(value: FeatureTensorType, index: Tensor,
+                 dim: int = 0) -> Tensor:
+    if isinstance(value, Tensor):
+        out: Optional[Tensor] = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we are in a background process, we write directly into a
+            # shared memory tensor to avoid an extra copy:
+            size = list(value.shape)
+            size[dim] = index.numel()
+            numel = math.prod(size)
+            storage = value.storage()._new_shared(numel)
+            out = value.new(storage).view(size)
+
+        return torch.index_select(value, dim, index, out=out)
+
+    elif isinstance(value, np.ndarray):
+        return torch.from_numpy(np.take(value, index, axis=dim))
+
+    raise ValueError(f"Encountered invalid feature tensor type "
+                     f"(got '{type(value)}')")
 
 
 def filter_node_store_(store: NodeStorage, out_store: NodeStorage,
@@ -35,7 +50,10 @@ def filter_node_store_(store: NodeStorage, out_store: NodeStorage,
             out_store.num_nodes = index.numel()
 
         elif store.is_node_attr(key):
-            index = index.to(value.device)
+            if isinstance(value, Tensor):
+                index = index.to(value.device)
+            elif isinstance(value, np.ndarray):
+                index = index.cpu()
             dim = store._parent().__cat_dim__(key, value, store)
             out_store[key] = index_select(value, index, dim=dim)
 
@@ -69,12 +87,17 @@ def filter_edge_store_(store: EdgeStorage, out_store: EdgeStorage, row: Tensor,
 
         elif store.is_edge_attr(key):
             dim = store._parent().__cat_dim__(key, value, store)
-            if perm is None:
+            if isinstance(value, Tensor):
                 index = index.to(value.device)
+            elif isinstance(value, np.ndarray):
+                index = index.cpu()
+            if perm is None:
                 out_store[key] = index_select(value, index, dim=dim)
             else:
-                perm = perm.to(value.device)
-                index = index.to(value.device)
+                if isinstance(value, Tensor):
+                    perm = perm.to(value.device)
+                elif isinstance(value, np.ndarray):
+                    perm = perm.cpu()
                 out_store[key] = index_select(value, perm[index], dim=dim)
 
     return store