#!/usr/bin/env python3 """ GraphSAINT-style Sampler for PyTorch Geometric. Mini-batch sampler for large graphs (140k+ nodes) using neighbor sampling. Compatible with base PyG installation (no torch-sparse or pyg-lib required). Usage: from models.spatiotemporal_gcn.sampler import GraphSAINTSampler sampler = GraphSAINTSampler( data=data, batch_size=256, num_neighbors=[256, 128, 64] ) """ import torch from torch.utils.data import DataLoader, Dataset from torch_geometric.data import Data from torch_geometric.utils import subgraph class GraphSAINTDataset(Dataset): """Dataset that samples node indices for mini-batching.""" def __init__(self, num_nodes: int, num_steps: int = 10): self.num_nodes = num_nodes self.num_steps = num_steps def __len__(self): return self.num_steps def __getitem__(self, idx): return torch.randint(0, self.num_nodes, (1,)) class GraphSAINTSampler: """ GraphSAINT-style mini-batch sampler for large graphs. Implements neighbor sampling to create subgraphs that fit in GPU memory. For each batch, samples seed nodes and their multi-hop neighbors. Args: data: Full graph with edge_index and node features. batch_size: Seed nodes per batch (default: 256). num_neighbors: Neighbors per layer [layer0, layer1, ...]. Default: [256, 128, 64] for 3-layer GCN. num_steps: Batches per epoch (default: 10). """ def __init__( self, data: Data, batch_size: int = 256, num_neighbors: list = None, num_steps: int = 10, ): if num_neighbors is None: num_neighbors = [256, 128, 64] self.data = data self.batch_size = batch_size self.num_neighbors = num_neighbors self.num_steps = num_steps self.num_nodes = data.num_nodes self.edge_index = data.edge_index if data.num_nodes > 100000: print(f"Sampler for large graph: {data.num_nodes:,} nodes") print(f" Batch size: {batch_size}") print(f" Layer depths: {num_neighbors}") def _sample_neighbors(self, seed_nodes: torch.Tensor) -> torch.Tensor: """ Sample multi-hop neighbors for seed nodes. Args: seed_nodes: Initial node indices. Returns: All sampled node indices (seed + neighbors). """ sampled = seed_nodes.unique() for num_neighbors in self.num_neighbors: if len(sampled) == 0: break mask = torch.isin(self.edge_index[0], sampled) neighbor_edges = self.edge_index[:, mask] if neighbor_edges.shape[1] == 0: break neighbors = neighbor_edges[1] if len(neighbors) > num_neighbors: neighbors = neighbors[torch.randperm(len(neighbors))[:num_neighbors]] sampled = torch.cat([sampled, neighbors]).unique() return sampled def _create_subgraph(self, node_indices: torch.Tensor) -> Data: edge_index, _, edge_mask = subgraph( node_indices, self.edge_index, relabel_nodes=True, return_edge_mask=True, ) subgraph_data = Data( x=self.data.x[node_indices], edge_index=edge_index, n_id=node_indices, ) if hasattr(self.data, 'y') and self.data.y is not None: subgraph_data.y = self.data.y[node_indices] return subgraph_data def __iter__(self): for _ in range(self.num_steps): seed_nodes = torch.randint(0, self.num_nodes, (self.batch_size,)) sampled_nodes = self._sample_neighbors(seed_nodes) batch = self._create_subgraph(sampled_nodes) yield batch def __len__(self): return self.num_steps class GraphSAINTConfig: """Configuration for GraphSAINT-style sampling.""" def __init__( self, batch_size: int = 256, num_neighbors: list = None, num_steps: int = 10, ): self.batch_size = batch_size self.num_neighbors = num_neighbors if num_neighbors is not None else [256, 128, 64] self.num_steps = num_steps def __repr__(self): return ( f"GraphSAINTConfig(\n" f" batch_size={self.batch_size},\n" f" num_neighbors={self.num_neighbors},\n" f" num_steps={self.num_steps}\n" f")" ) def create_graph_saint_loader( data: Data, batch_size: int = 256, num_neighbors: list = None, num_steps: int = 10, ): """ Create a GraphSAINT-style sampler for large graph training. Args: data: Full graph data with edge_index and features. batch_size: Seed nodes per batch (default: 256). num_neighbors: Layer-wise neighbor counts (default: [256, 128, 64]). num_steps: Batches per epoch (default: 10). Returns: GraphSAINTSampler: Mini-batch iterator. """ return GraphSAINTSampler( data=data, batch_size=batch_size, num_neighbors=num_neighbors, num_steps=num_steps, ) def main(): """Example usage with dummy data.""" print("=" * 60) print("GraphSAINT-style Sampler Demo") print("=" * 60) print("\nCreating dummy graph (10k nodes)...") N = 10000 num_features = 48 edge_index = torch.randint(0, N, (2, N * 3)) x = torch.randn(N, num_features) y = torch.randint(0, 3, (N,)) data = Data(x=x, y=y, edge_index=edge_index) print(f" Nodes: {data.num_nodes:,}") print(f" Edges: {data.num_edges:,}") print(f" Features: {data.num_node_features}") print("\nCreating sampler...") config = GraphSAINTConfig( batch_size=256, num_neighbors=[256, 128, 64], num_steps=5, ) print(config) loader = create_graph_saint_loader( data=data, batch_size=config.batch_size, num_neighbors=config.num_neighbors, num_steps=config.num_steps, ) print(f"\nIterating through {len(loader)} batches...") for i, batch in enumerate(loader): print(f" Batch {i+1}/{len(loader)}:") print(f" Nodes: {batch.num_nodes:,}") print(f" Edges: {batch.num_edges:,}") print(f" Features: {batch.x.shape}") print(f" Node IDs: {batch.n_id.shape}") if i >= 2: break print("\n" + "=" * 60) print("Sampler ready for training!") print("=" * 60) print("\nFor your 140k node graph:") print(" 1. Load graph: data = load_your_graph()") print(" 2. Create loader: loader = create_graph_saint_loader(data, batch_size=256)") print(" 3. Train: for batch in loader: out = model(batch.x, batch.edge_index)") print("\nRecommended for 4GB GPU:") print(" - batch_size: 256") print(" - num_neighbors: [256, 128, 64]") if __name__ == '__main__': main()