cugraph_pyg.loader.NeighborLoader#

class cugraph_pyg.loader.NeighborLoader(

data: torch_geometric.data.Data | torch_geometric.data.HeteroData | Tuple[torch_geometric.data.FeatureStore, torch_geometric.data.GraphStore],

num_neighbors: List[int] | Dict[torch_geometric.typing.EdgeType, List[int]],

input_nodes: torch_geometric.typing.InputNodes = None,

input_time: torch_geometric.typing.OptTensor = None,

replace: bool = False,

subgraph_type: torch_geometric.typing.SubgraphType | str = 'directional',

disjoint: bool = False,

temporal_strategy: str = 'uniform',

time_attr: str | None = None,

weight_attr: str | None = None,

transform: Callable | None = None,

transform_sampler_output: Callable | None = None,

is_sorted: bool = False,

filter_per_worker: bool | None = None,

neighbor_sampler: torch_geometric.sampler.NeighborSampler | None = None,

directed: bool = True,

batch_size: int = 16,

compression: str | None = None,

local_seeds_per_call: int | None = None,

temporal_comparison: str | None = None,

input_start_time: torch_geometric.typing.OptTensor = None,

input_end_time: torch_geometric.typing.OptTensor = None,

**kwargs,

)[source]#

Duck-typed version of torch_geometric.loader.NeighborLoader

Node loader that implements the neighbor sampling algorithm used in GraphSAGE.

__init__(

data: torch_geometric.data.Data | torch_geometric.data.HeteroData | Tuple[torch_geometric.data.FeatureStore, torch_geometric.data.GraphStore],

num_neighbors: List[int] | Dict[torch_geometric.typing.EdgeType, List[int]],

input_nodes: torch_geometric.typing.InputNodes = None,

input_time: torch_geometric.typing.OptTensor = None,

replace: bool = False,

subgraph_type: torch_geometric.typing.SubgraphType | str = 'directional',

disjoint: bool = False,

temporal_strategy: str = 'uniform',

time_attr: str | None = None,

weight_attr: str | None = None,

transform: Callable | None = None,

transform_sampler_output: Callable | None = None,

is_sorted: bool = False,

filter_per_worker: bool | None = None,

neighbor_sampler: torch_geometric.sampler.NeighborSampler | None = None,

directed: bool = True,

batch_size: int = 16,

compression: str | None = None,

local_seeds_per_call: int | None = None,

temporal_comparison: str | None = None,

input_start_time: torch_geometric.typing.OptTensor = None,

input_end_time: torch_geometric.typing.OptTensor = None,

**kwargs,

)[source]#

data: Data, HeteroData, or Tuple[FeatureStore, GraphStore]: See torch_geometric.loader.NeighborLoader.
num_neighbors: List[int] or Dict[EdgeType, List[int]]: Fanout values. See torch_geometric.loader.NeighborLoader.
input_nodes: InputNodes: Input nodes for sampling. See torch_geometric.loader.NeighborLoader.
input_time: OptTensor (optional): See torch_geometric.loader.NeighborLoader.
replace: bool (optional, default=False): Whether to sample with replacement. See torch_geometric.loader.NeighborLoader.
subgraph_type: Union[SubgraphType, str] (optional, default=’directional’): The type of subgraph to return. Currently only ‘directional’ is supported. See torch_geometric.loader.NeighborLoader.
disjoint: bool (optional, default=False): Whether to perform disjoint sampling. See torch_geometric.loader.NeighborLoader.
temporal_strategy: str (optional, default=’uniform’): The temporal sampling strategy ('uniform' or 'last'). 'last' selects the most recent neighbors within each seed’s fixed time window and therefore requires input_start_time and input_end_time. See torch_geometric.loader.NeighborLoader.
time_attr: str (optional, default=None): Used for temporal sampling. See torch_geometric.loader.NeighborLoader.
weight_attr: str (optional, default=None): Used for biased sampling. See torch_geometric.loader.NeighborLoader.
transform: Callable (optional, default=None): See torch_geometric.loader.NeighborLoader.
transform_sampler_output: Callable (optional, default=None): See torch_geometric.loader.NeighborLoader.
is_sorted: bool (optional, default=False): Ignored by cuGraph. See torch_geometric.loader.NeighborLoader.
filter_per_worker: bool (optional, default=False): Currently ignored by cuGraph, but this may change once in-memory sampling is implemented. See torch_geometric.loader.NeighborLoader.
neighbor_sampler: torch_geometric.sampler.NeighborSampler: (optional, default=None) Not supported by cuGraph. See torch_geometric.loader.NeighborLoader.
directed: bool (optional, default=True): Deprecated. See torch_geometric.loader.NeighborLoader.
batch_size: int (optional, default=16): The number of input nodes per output minibatch. See torch.utils.dataloader.
compression: str (optional, default=None): The compression type to use if writing samples to disk. If not provided, it is automatically chosen.
local_seeds_per_call: int (optional, default=None): The number of seeds to process within a single sampling call. Manually tuning this parameter is not recommended but reducing it may conserve GPU memory. The total number of seeds processed per sampling call is equal to the sum of this parameter across all workers. If not provided, it will be automatically calculated. See cugraph_pyg.sampler.BaseDistributedSampler.
temporal_comparison: str (optional, default=’monotonically_decreasing’): The comparison operator for temporal sampling (‘strictly_increasing’, ‘monotonically_increasing’, ‘strictly_decreasing’, ‘monotonically_decreasing’). Fixed-window sampling (input_start_time / input_end_time) only supports 'monotonically_increasing'; passing any other value raises a ValueError. temporal_strategy='last' also requires 'monotonically_increasing' and sets it automatically. See cugraph_pyg.sampler.BaseDistributedSampler.
input_start_time: OptTensor (optional): Lower time-window bounds for each input node. Must be passed together with input_end_time and instead of input_time.
input_end_time: OptTensor (optional): Upper time-window bounds for each input node. Must be passed together with input_start_time and instead of input_time.
**kwargs: Other keyword arguments passed to the superclass.

Methods

__init__(data, num_neighbors[, input_nodes, ...])

data: Data, HeteroData, or Tuple[FeatureStore, GraphStore]