Akiba So
fc468464b2
Context: Build a spatial risk assessment system correlating air quality data with children's respiratory disease incidence across Wuhan. Approach: FastAPI backend serving PostGIS spatial queries, React frontend with Deck.gl maps, and a PyTorch SpatialTemporalGCN pipeline for multi-day (1d/3d/7d) risk prediction. Changes: - backend/ — FastAPI API with auth (JWT), alerts, risk analysis, geocoded case data, grid statistics, and report endpoints - frontend/ — React dashboard with interactive risk maps, alert monitoring, district comparison charts, and timeline player - models/ — SpatialTemporalGCN model with trained weights and ONNX export for inference - scripts/ — ETL pipeline for weather + medical data, grid generation, feature engineering, training, and daily inference - deploy/ — Docker Compose configs for backend, frontend, and MLflow - docs/ — API docs, deployment guide, user guide, and code review Impact: Enables spatial risk visualization, alert monitoring, and ML-driven health risk forecasting for environmental health teams.
245 lines
7.1 KiB
Python
245 lines
7.1 KiB
Python
#!/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()
|