k3-node 1.0.0__py3-none-any.whl

k3_node/loader/neighbor_sampler.py

@@ -0,0 +1,167 @@
+from typing import Any, Callable, List, NamedTuple, Optional, Tuple, Union
+
+import numpy as np
+
+try:
+    import torch
+    from torch import Tensor
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    Tensor = type(None)
+    BaseDataLoader = object
+
+from k3_node.loader.sampler_utils import FastGraph
+
+
+class EdgeIndex(NamedTuple):
+    edge_index: Any
+    e_id: Optional[Any]
+    size: Tuple[int, int]
+
+    def to(self, *args, **kwargs):
+        if hasattr(self.edge_index, 'to'):
+            edge_index = self.edge_index.to(*args, **kwargs)
+            e_id = self.e_id.to(*args, **kwargs) if self.e_id is not None and hasattr(self.e_id, 'to') else None
+            return EdgeIndex(edge_index, e_id, self.size)
+        return self
+
+
+class Adj(NamedTuple):
+    adj_t: Any
+    e_id: Optional[Any]
+    size: Tuple[int, int]
+
+    def to(self, *args, **kwargs):
+        if hasattr(self.adj_t, 'to'):
+            adj_t = self.adj_t.to(*args, **kwargs)
+            e_id = self.e_id.to(*args, **kwargs) if self.e_id is not None and hasattr(self.e_id, 'to') else None
+            return Adj(adj_t, e_id, self.size)
+        return self
+
+
+class NeighborSampler(BaseDataLoader):
+    r"""The legacy layer-by-layer bipartite neighbor sampler from "Inductive Representation
+    Learning on Large Graphs".
+
+    Args:
+        edge_index (Tensor): Edge connectivity.
+        sizes (List[int]): Number of neighbors to sample per layer.
+        node_idx (Tensor, optional): Seed nodes to consider. (default: :obj:`None`)
+        num_nodes (int, optional): Number of nodes. (default: :obj:`None`)
+        return_e_id (bool, optional): Whether to return edge IDs. (default: :obj:`True`)
+        transform (callable, optional): Transform function. (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        edge_index: Any,
+        sizes: List[int],
+        node_idx: Optional[Any] = None,
+        num_nodes: Optional[int] = None,
+        return_e_id: bool = True,
+        transform: Optional[Callable] = None,
+        **kwargs,
+    ):
+        kwargs.pop('dataset', None)
+        kwargs.pop('collate_fn', None)
+
+        self.sizes = sizes
+        self.return_e_id = return_e_id
+        self.transform = transform
+
+        is_torch = torch is not None and isinstance(edge_index, Tensor)
+        if is_torch:
+            np_edge_index = edge_index.detach().cpu().numpy()
+        else:
+            np_edge_index = np.asarray(edge_index)
+
+        if num_nodes is None:
+            num_nodes = int(np.max(np_edge_index) + 1) if np_edge_index.size > 0 else 0
+
+        self.num_nodes = num_nodes
+        self.graph = FastGraph(np_edge_index, num_nodes=num_nodes)
+
+        if node_idx is None:
+            node_idx = torch.arange(num_nodes) if is_torch else np.arange(num_nodes)
+        elif is_torch and isinstance(node_idx, Tensor) and node_idx.dtype == torch.bool:
+            node_idx = node_idx.nonzero(as_tuple=False).view(-1)
+        elif isinstance(node_idx, np.ndarray) and node_idx.dtype == bool:
+            node_idx = np.nonzero(node_idx)[0]
+
+        self.node_idx = node_idx
+        idx_list = node_idx.tolist() if hasattr(node_idx, 'tolist') else list(node_idx)
+
+        if torch is not None:
+            super().__init__(idx_list, collate_fn=self.sample, **kwargs)
+        else:
+            self.dataset = idx_list
+            self.collate_fn = self.sample
+
+    def sample(self, batch: Any) -> Any:
+        is_torch = torch is not None
+        if is_torch and not isinstance(batch, Tensor):
+            batch = torch.tensor(batch, dtype=torch.long)
+            batch_size = batch.numel()
+            curr_nodes = batch.tolist()
+        else:
+            batch_size = len(batch)
+            curr_nodes = list(batch)
+
+        adjs = []
+        n_id = list(curr_nodes)
+        node_map = {n: i for i, n in enumerate(n_id)}
+
+        for size in self.sizes:
+            rows, cols, e_ids = [], [], []
+            targets = list(curr_nodes)
+            next_curr_nodes = []
+
+            for target in targets:
+                srcs, edges = self.graph.get_neighbors(target)
+                count = len(srcs)
+                if count == 0:
+                    continue
+
+                if size == -1 or size >= count:
+                    chosen = np.arange(count)
+                else:
+                    chosen = np.random.choice(count, size=size, replace=False)
+
+                for s, e in zip(srcs[chosen], edges[chosen]):
+                    s = int(s)
+                    e = int(e)
+                    if s not in node_map:
+                        node_map[s] = len(n_id)
+                        n_id.append(s)
+                        next_curr_nodes.append(s)
+
+                    rows.append(node_map[s])
+                    cols.append(node_map[target])
+                    e_ids.append(e)
+
+            curr_nodes = list(n_id)
+
+            if len(rows) > 0:
+                edge_index = np.stack([np.array(rows, dtype=np.int64), np.array(cols, dtype=np.int64)], axis=0)
+                edge_ids = np.array(e_ids, dtype=np.int64) if self.return_e_id else None
+            else:
+                edge_index = np.empty((2, 0), dtype=np.int64)
+                edge_ids = np.empty(0, dtype=np.int64) if self.return_e_id else None
+
+            bipartite_size = (len(n_id), len(targets))
+
+            if is_torch:
+                edge_index = torch.from_numpy(edge_index)
+                if edge_ids is not None:
+                    edge_ids = torch.from_numpy(edge_ids)
+
+            adjs.append(EdgeIndex(edge_index, edge_ids, bipartite_size))
+
+        adjs = adjs[0] if len(adjs) == 1 else adjs[::-1]
+        out_n_id = torch.tensor(n_id, dtype=torch.long) if is_torch else np.array(n_id, dtype=np.int64)
+        out = (batch_size, out_n_id, adjs)
+        return self.transform(*out) if self.transform is not None else out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(sizes={self.sizes})'

k3_node/loader/node_loader.py

@@ -0,0 +1,185 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+
+try:
+    import torch
+    from torch import Tensor
+except ImportError:
+    torch = None
+    Tensor = type(None)
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader.base import BaseDataLoader, DataLoaderIterator
+from k3_node.loader.mixin import AffinityMixin, LogMemoryMixin, MultithreadingMixin
+from k3_node.loader.utils import (
+    filter_data,
+    filter_hetero_data,
+    get_input_nodes,
+    infer_filter_per_worker,
+)
+from k3_node.loader.keras_dataset import loader_bases
+
+
+@dataclass
+class NodeSamplerInput:
+    input_id: Optional[Any]
+    node: Any
+    time: Optional[Any] = None
+    input_type: Optional[str] = None
+
+    def __getitem__(self, index: Any) -> 'NodeSamplerInput':
+        if torch is not None and not isinstance(index, Tensor):
+            index = torch.as_tensor(index, dtype=torch.long)
+        return NodeSamplerInput(
+            input_id=self.input_id[index] if self.input_id is not None else index,
+            node=self.node[index],
+            time=self.time[index] if self.time is not None else None,
+            input_type=self.input_type,
+        )
+
+
+@dataclass
+class SamplerOutput:
+    node: Any
+    row: Any
+    col: Any
+    edge: Optional[Any] = None
+    batch: Optional[Any] = None
+    num_sampled_nodes: Optional[List[int]] = None
+    num_sampled_edges: Optional[List[int]] = None
+    orig_row: Optional[Any] = None
+    orig_col: Optional[Any] = None
+    metadata: Optional[Any] = None
+
+
+@dataclass
+class HeteroSamplerOutput:
+    node: Dict[str, Any]
+    row: Dict[Tuple[str, str, str], Any]
+    col: Dict[Tuple[str, str, str], Any]
+    edge: Dict[Tuple[str, str, str], Optional[Any]]
+    batch: Optional[Dict[str, Any]] = None
+    num_sampled_nodes: Optional[Dict[str, List[int]]] = None
+    num_sampled_edges: Optional[Dict[Tuple[str, str, str], List[int]]] = None
+    orig_row: Optional[Dict[Tuple[str, str, str], Any]] = None
+    orig_col: Optional[Dict[Tuple[str, str, str], Any]] = None
+    metadata: Optional[Any] = None
+
+
+class NodeLoader(*loader_bases(BaseDataLoader), AffinityMixin, MultithreadingMixin, LogMemoryMixin):
+    r"""A data loader that performs mini-batch sampling from node information."""
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        node_sampler: Any,
+        input_nodes: Any = None,
+        input_time: Optional[Any] = None,
+        transform: Optional[Callable] = None,
+        transform_sampler_output: Optional[Callable] = None,
+        filter_per_worker: Optional[bool] = None,
+        custom_cls: Optional[Any] = None,
+        input_id: Optional[Any] = None,
+        **kwargs,
+    ):
+        if filter_per_worker is None:
+            filter_per_worker = infer_filter_per_worker(data)
+
+        self.data = data
+        self.node_sampler = node_sampler
+        self.input_nodes = input_nodes
+        self.input_time = input_time
+        self.transform = transform
+        self.transform_sampler_output = transform_sampler_output
+        self.filter_per_worker = filter_per_worker
+        self.custom_cls = custom_cls
+        self.input_id = input_id
+
+        kwargs.pop('dataset', None)
+        kwargs.pop('collate_fn', None)
+
+        input_type, input_nodes, input_id = get_input_nodes(data, input_nodes, input_id)
+
+        self.input_data = NodeSamplerInput(
+            input_id=input_id,
+            node=input_nodes,
+            time=input_time,
+            input_type=input_type,
+        )
+
+        num_inputs = input_nodes.size(0) if hasattr(input_nodes, 'size') else len(input_nodes)
+        iterator = range(num_inputs)
+
+        if torch is not None:
+            super().__init__(iterator, collate_fn=self.collate_fn, **kwargs)
+        else:
+            self.dataset = iterator
+            self.collate_fn = self.collate_fn
+
+    def __call__(self, index: Any) -> Union[Data, HeteroData]:
+        out = self.collate_fn(index)
+        if not self.filter_per_worker:
+            out = self.filter_fn(out)
+        return out
+
+    def collate_fn(self, index: Any) -> Any:
+        input_data = self.input_data[index]
+        out = self.node_sampler.sample_from_nodes(input_data)
+        if self.filter_per_worker:
+            out = self.filter_fn(out)
+        return out
+
+    def filter_fn(self, out: Any) -> Union[Data, HeteroData]:
+        if self.transform_sampler_output:
+            out = self.transform_sampler_output(out)
+
+        if isinstance(out, SamplerOutput):
+            perm = getattr(self.node_sampler, 'edge_permutation', None)
+            data = filter_data(self.data, out.node, out.row, out.col, out.edge, perm)
+
+            data.n_id = out.node
+            if out.edge is not None:
+                data.e_id = out.edge
+            data.batch = out.batch
+            data.num_sampled_nodes = out.num_sampled_nodes
+            data.num_sampled_edges = out.num_sampled_edges
+
+            meta = out.metadata or (out.node, None)
+            data.input_id = meta[0]
+            data.batch_size = meta[0].size(0) if hasattr(meta[0], 'size') else len(meta[0])
+
+        elif isinstance(out, HeteroSamplerOutput):
+            perm = getattr(self.node_sampler, 'edge_permutation', None)
+            data = filter_hetero_data(self.data, out.node, out.row, out.col, out.edge, perm)
+
+            for key, node in out.node.items():
+                data[key].n_id = node
+
+            for key, edge in (out.edge or {}).items():
+                if edge is not None:
+                    data[key].e_id = edge
+
+            if out.batch is not None:
+                data.set_value_dict('batch', out.batch)
+            if out.num_sampled_nodes is not None:
+                data.set_value_dict('num_sampled_nodes', out.num_sampled_nodes)
+            if out.num_sampled_edges is not None:
+                data.set_value_dict('num_sampled_edges', out.num_sampled_edges)
+
+            input_type = self.input_data.input_type
+            meta = out.metadata or (out.node.get(input_type, None), None)
+            data[input_type].input_id = meta[0]
+            data[input_type].batch_size = meta[0].size(0) if hasattr(meta[0], 'size') else len(meta[0])
+
+        else:
+            data = out
+
+        return data if self.transform is None else self.transform(data)
+
+    def _get_iterator(self) -> Iterator:
+        if self.filter_per_worker:
+            return super()._get_iterator()
+        return DataLoaderIterator(super()._get_iterator(), self.filter_fn)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'
+

k3_node/loader/prefetch.py

@@ -0,0 +1,115 @@
+import warnings
+from contextlib import nullcontext
+from functools import partial
+from typing import Any, Optional
+
+try:
+    import torch
+except ImportError:
+    torch = None
+    DataLoader = object
+
+
+class DeviceHelper:
+    def __init__(self, device: Optional[Any] = None):
+        if torch is None:
+            self.device = 'cpu'
+            self.is_gpu = False
+            self.stream = None
+            self.stream_context = nullcontext
+            self.module = None
+            return
+
+        with_cuda = torch.cuda.is_available()
+        with_xpu = hasattr(torch, 'xpu') and torch.xpu.is_available()
+
+        if device is None:
+            if with_cuda:
+                device = 'cuda'
+            elif with_xpu:
+                device = 'xpu'
+            else:
+                device = 'cpu'
+
+        self.device = torch.device(device)
+        self.is_gpu = self.device.type in ['cuda', 'xpu']
+
+        if ((self.device.type == 'cuda' and not with_cuda)
+                or (self.device.type == 'xpu' and not with_xpu)):
+            warnings.warn(
+                f"Requested device '{self.device.type}' is not available, falling back to CPU",
+                stacklevel=2,
+            )
+            self.device = torch.device('cpu')
+            self.is_gpu = False
+
+        self.stream = None
+        self.stream_context = nullcontext
+        self.module = getattr(torch, self.device.type) if self.is_gpu else None
+
+    def maybe_init_stream(self) -> None:
+        if self.is_gpu and self.module is not None:
+            self.stream = self.module.Stream()
+            self.stream_context = partial(self.module.stream, stream=self.stream)
+
+    def maybe_wait_stream(self) -> None:
+        if self.stream is not None and self.module is not None:
+            self.module.current_stream().wait_stream(self.stream)
+
+
+class PrefetchLoader:
+    r"""A prefetcher class for asynchronously transferring data of a DataLoader
+    from host memory to device memory.
+
+    Args:
+        loader (DataLoader): The data loader.
+        device (torch.device, optional): The device to load the data to. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        loader: Any,
+        device: Optional[Any] = None,
+    ):
+        self.loader = loader
+        self.device_helper = DeviceHelper(device)
+
+    def non_blocking_transfer(self, batch: Any) -> Any:
+        if not self.device_helper.is_gpu:
+            return batch
+        if isinstance(batch, (list, tuple)):
+            return type(batch)(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()}
+
+        if hasattr(batch, 'pin_memory'):
+            batch = batch.pin_memory()
+        if hasattr(batch, 'to'):
+            return batch.to(self.device_helper.device, non_blocking=True)
+        return batch
+
+    def __iter__(self) -> Any:
+        first = True
+        self.device_helper.maybe_init_stream()
+
+        batch = None
+        for next_batch in self.loader:
+            with self.device_helper.stream_context():
+                next_batch = self.non_blocking_transfer(next_batch)
+
+            if not first:
+                yield batch
+            else:
+                first = False
+
+            self.device_helper.maybe_wait_stream()
+            batch = next_batch
+
+        if batch is not None:
+            yield batch
+
+    def __len__(self) -> int:
+        return len(self.loader)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.loader})'
+

k3_node/loader/random_node_loader.py

@@ -0,0 +1,89 @@
+import math
+from typing import Union
+
+import numpy as np
+
+try:
+    import torch
+    import torch.utils.data
+    from torch import Tensor
+    BaseDataLoader = torch.utils.data.DataLoader
+except ImportError:
+    torch = None
+    Tensor = type(None)
+    BaseDataLoader = object
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader.keras_dataset import loader_bases
+
+
+class RandomNodeLoader(*loader_bases(BaseDataLoader)):
+    r"""A data loader that randomly samples nodes within a graph and returns
+    their induced subgraph.
+
+    Args:
+        data (Data or HeteroData): The graph data object.
+        num_parts (int): The number of partitions.
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        num_parts: int,
+        **kwargs,
+    ):
+        self.data = data
+        self.num_parts = num_parts
+
+        if isinstance(data, HeteroData):
+            node_dict = {}
+            total = 0
+            for node_type in data.node_types:
+                count = data[node_type].num_nodes
+                node_dict[node_type] = (total, total + count)
+                total += count
+            self.node_dict = node_dict
+            self.num_nodes = total
+        else:
+            self.edge_index = data.edge_index
+            self.num_nodes = data.num_nodes
+
+        kwargs.pop('dataset', None)
+        kwargs.pop('batch_size', None)
+        kwargs.pop('collate_fn', None)
+
+        batch_size = math.ceil(self.num_nodes / num_parts) if num_parts > 0 else self.num_nodes
+
+        if torch is not None:
+            super().__init__(
+                range(self.num_nodes),
+                batch_size=batch_size,
+                collate_fn=self.collate_fn,
+                **kwargs,
+            )
+        else:
+            self.dataset = range(self.num_nodes)
+            self.batch_size = batch_size
+            self.collate_fn = self.collate_fn
+
+    def collate_fn(self, index):
+        if torch is not None and not isinstance(index, Tensor):
+            index = torch.tensor(index, dtype=torch.long)
+        elif torch is None:
+            index = np.asarray(index, dtype=np.int64)
+
+        if isinstance(self.data, Data):
+            return self.data.subgraph(index)
+
+        elif isinstance(self.data, HeteroData):
+            node_dict = {}
+            for key, (start, end) in self.node_dict.items():
+                if torch is not None and isinstance(index, Tensor):
+                    mask = (index >= start) & (index < end)
+                    node_dict[key] = index[mask] - start
+                else:
+                    idx_np = np.asarray(index)
+                    mask = (idx_np >= start) & (idx_np < end)
+                    node_dict[key] = idx_np[mask] - start
+            return self.data.subgraph(node_dict)
+