torch_geometric.sampler

BaseSampler

A base class that initializes a graph sampler and provides sample_from_nodes() and sample_from_edges() routines.

NeighborSampler

An implementation of an in-memory (heterogeneous) neighbor sampler used by NeighborLoader.

HGTSampler

An implementation of an in-memory heterogeneous layer-wise sampler user by HGTLoader.

class BaseSampler[source]

A base class that initializes a graph sampler and provides sample_from_nodes() and sample_from_edges() routines.

Note

Any data stored in the sampler will be replicated across data loading workers that use the sampler since each data loading worker holds its own instance of a sampler. As such, it is recommended to limit the amount of information stored in the sampler.

abstract sample_from_nodes(index: Tuple[Tensor, Tensor, Optional[Tensor]], **kwargs) Union[HeteroSamplerOutput, SamplerOutput][source]

Performs sampling from the nodes specified in index, returning a sampled subgraph in the specified output format.

Parameters

index (Tensor) – The node indices to start sampling from.

abstract sample_from_edges(index: Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor]], **kwargs) Union[HeteroSamplerOutput, SamplerOutput][source]

Performs sampling from the edges specified in index, returning a sampled subgraph in the specified output format.

Parameters

index (Tuple[Tensor, Tensor, Tensor, Optional[Tensor]]) – The (1) source node indices, the (2) destination node indices, the (3) edge labels and the (4) optional timestamp of edges to start sampling from.

property edge_permutation: Union[Tensor, None, Dict[Tuple[str, str, str], Optional[Tensor]]]

If the sampler performs any modification of edge ordering in the original graph, this function is expected to return the permutation tensor that defines the permutation from the edges in the original graph and the edges used in the sampler. If no such permutation was applied, None is returned. For heterogeneous graphs, the expected return type is a permutation tensor for each edge type.

class NeighborSampler(data: Union[Data, HeteroData, Tuple[FeatureStore, GraphStore]], num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]], replace: bool = False, directed: bool = True, disjoint: bool = False, temporal_strategy: str = 'uniform', input_type: Optional[Any] = None, time_attr: Optional[str] = None, is_sorted: bool = False, share_memory: bool = False)[source]

An implementation of an in-memory (heterogeneous) neighbor sampler used by NeighborLoader.

property is_temporal: bool
property disjoint: bool
class HGTSampler(data: HeteroData, num_samples: Union[List[int], Dict[str, List[int]]], input_type: Optional[Any] = None, is_sorted: bool = False, share_memory: bool = False)[source]

An implementation of an in-memory heterogeneous layer-wise sampler user by HGTLoader.