torch_geometric.sampler
|
A base class that initializes a graph sampler and provides |
An implementation of an in-memory (heterogeneous) neighbor sampler used by |
|
An implementation of an in-memory heterogeneous layer-wise sampler user by |
- class BaseSampler[source]
A base class that initializes a graph sampler and provides
sample_from_nodes()
andsample_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
.