从交通拥堵到疫情预测：手把手教你用STGNN模型解决5个城市计算难题
从交通拥堵到疫情预测：STGNN模型实战指南
城市计算领域正迎来一场由时空图神经网络（STGNN）驱动的技术变革。这种能够同时捕捉空间关联与时间动态的AI模型，正在重塑我们对城市复杂系统的理解方式。不同于传统时序预测方法，STGNN通过图结构建模区域间的多维关系，在交通管理、公共安全、环境监测等场景展现出惊人的预测精度。本文将带您深入五个典型应用场景，手把手构建端到端的预测解决方案。
1. 短时交通流预测实战
PeMS交通数据集已成为评估STGNN性能的黄金标准。这个覆盖加州主要高速公路的传感器网络，每30秒记录一次交通流量、速度和占有率。我们首先需要解决的关键问题是：如何将原始传感器数据转化为STGNN可处理的时空图？
1.1 数据预处理与图构建
原始PeMS数据通常以CSV格式存储，包含三组关键字段：
- 静态属性：传感器ID、经纬度坐标、车道数
- 动态特征：时间戳、流量(veh/h)、速度(mph)、占有率(%)
- 拓扑信息：传感器间的连接关系（基于实际路网）
import pandas as pd import numpy as np # 加载并清洗数据 def load_pems_data(file_path): df = pd.read_csv(file_path) # 处理缺失值：线性插值+前后填充 df = df.interpolate().fillna(method='bfill').fillna(method='ffill') # 标准化：各传感器独立归一化 grouped = df.groupby('sensor_id') df['flow'] = grouped['flow'].transform(lambda x: (x - x.mean())/x.std()) return df # 构建基于距离的邻接矩阵 def build_adj_matrix(coords_file, threshold=5): coords = pd.read_csv(coords_file) dist_matrix = np.zeros((len(coords), len(coords))) for i, (x1, y1) in enumerate(zip(coords['longitude'], coords['latitude'])): for j, (x2, y2) in enumerate(zip(coords['longitude'], coords['latitude'])): dist = haversine(x1, y1, x2, y2) # 高斯核函数加权 dist_matrix[i,j] = np.exp(-dist**2/threshold**2) if dist <= threshold*2 else 0 return dist_matrix提示：实际应用中建议同时构建多种图结构（基于距离、路网拓扑、交通流相关性），通过实验选择最优组合或采用多图融合架构。
1.2 模型选型与调优
STGCN和DCRNN是交通预测中最成熟的两种架构，其性能对比见下表：
| 模型 | 参数量 | 训练速度 | MAE(5min) | RMSE(15min) | 适用场景 |
|---|---|---|---|---|---|
| STGCN | 约350K | 快(1.2x) | 2.81 | 5.32 | 算力有限时 |
| DCRNN | 约480K | 慢 | 2.63 | 4.97 | 高精度要求 |
| GraphWaveNet | 约520K | 中等 | 2.57 | 4.85 | 动态图关系 |
调优时重点关注三个超参数：
- 时间窗口大小：通常取12个时间步（1小时历史）
- 图卷积层数：超过3层可能导致过平滑
- 扩散步长：DCRNN中建议设置为2-3
# 使用PyTorch Geometric实现STGCN import torch from torch_geometric.nn import STConv class STGCN(torch.nn.Module): def __init__(self, num_nodes, in_channels): super().__init__() self.conv1 = STConv(num_nodes, in_channels, 32, kernel_size=3) self.conv2 = STConv(num_nodes, 32, 64, kernel_size=3) self.linear = torch.nn.Linear(64, 12) # 预测未来12个时间步 def forward(self, x, edge_index): x = F.relu(self.conv1(x, edge_index)) x = self.conv2(x, edge_index) return self.linear(x)2. 区域犯罪热点预警系统
纽约市犯罪数据报告系统记录了2010年以来的300多万条犯罪事件。我们将这些数据按社区划分，构建时空预测模型。
2.1 多源数据融合
犯罪模式受多种因素影响，需要整合：
- 核心数据：犯罪类型、时间、社区位置
- 环境因子：POI分布、人口密度、经济指标
- 时序特征：节假日、天气、特殊事件
# 构建多模态图结构 def build_crime_graph(): # 基于行政边界的社区图 admin_graph = nx.read_gml('nyc_community.gml') # 基于人口流动的交互图 mobility = pd.read_csv('citibike_flow.csv') flow_graph = nx.from_pandas_edgelist(mobility, 'source', 'target', 'count') # 基于POI相似性的语义图 poi_sim = cosine_similarity(poi_matrix) semantic_graph = nx.from_numpy_array(poi_sim) return [admin_graph, flow_graph, semantic_graph]2.2 异构图神经网络应用
使用RGCN处理不同类型的关系：
from torch_geometric.nn import RGCNConv class CrimePredictor(torch.nn.Module): def __init__(self, num_relations): super().__init__() self.conv1 = RGCNConv(64, 64, num_relations) self.conv2 = RGCNConv(64, 64, num_relations) self.temporal = nn.GRU(64, 64) def forward(self, x, edge_indices, edge_types): for i in range(len(edge_indices)): x = self.conv1(x, edge_indices[i], edge_types[i]) x, _ = self.temporal(x) return x注意：犯罪预测涉及伦理问题，建议：
- 仅输出区域级风险指数，不关联个人
- 加入公平性约束，防止算法偏见
- 结果需经人工复核方可作为决策依据
3. 共享单车需求调度优化
芝加哥Divvy单车系统每天产生约2万次骑行记录。精准预测各站点的车辆需求可降低运营成本30%以上。
3.1 动态图构建策略
单车需求具有明显的时空传播特性：
- 流入流出模式：早高峰从住宅区流向商业区
- 级联效应：一个站点的短缺会影响邻近站点
- 外部影响：天气、活动、地铁故障等
# 动态调整图结构 def update_graph(hourly_flow): adj = np.zeros((n_stations, n_stations)) for i in range(n_stations): for j in range(n_stations): # 流向强度+距离衰减 adj[i,j] = hourly_flow[i,j] * np.exp(-distance[i,j]/500) return normalize(adj)3.2 时空注意力机制
ASTGCN模型能自动学习不同时间尺度的影响：
class TemporalAttention(nn.Module): def __init__(self, in_dim): super().__init__() self.query = nn.Linear(in_dim, in_dim) self.key = nn.Linear(in_dim, in_dim) def forward(self, x): Q = self.query(x) K = self.key(x) scores = torch.matmul(Q, K.transpose(-2,-1)) / np.sqrt(in_dim) return torch.softmax(scores, dim=-1)实际部署时，建议采用以下优化策略：
- 增量训练：每周更新模型参数
- 弹性调度：预留5%车辆应对预测误差
- 激励机制：引导用户参与平衡调度
4. 空气质量站点插值技术
全国环境监测网约有1,500个站点，需要通过STGNN实现高分辨率污染分布图。
4.1 物理约束建模
将大气扩散方程融入模型：
def advection_diffusion(x, adj, wind): # x: 污染物浓度 [n_nodes, 1] # wind: 风速风向 [n_nodes, 2] laplacian = compute_laplacian(adj) grad = compute_gradient(wind) return -wind @ grad + 0.1 * laplacian @ x # 0.1为扩散系数4.2 多任务学习框架
同时预测多种污染物：
| 任务 | 评估指标 | 权重 |
|---|---|---|
| PM2.5 | MAE | 0.4 |
| O3 | MAPE | 0.3 |
| NO2 | RMSE | 0.3 |
class MultiTaskHead(nn.Module): def __init__(self, in_dim): super().__init__() self.shared_backbone = STGNN(in_dim) self.heads = nn.ModuleDict({ 'pm25': nn.Linear(64, 1), 'o3': nn.Linear(64, 1), 'no2': nn.Linear(64, 1) }) def forward(self, x, adj): shared = self.shared_backbone(x, adj) return {k: head(shared) for k, head in self.heads.items()}5. 流行病传播模拟系统
5.1 基于SIR的混合建模
将经典流行病学模型与STGNN结合：
class SIRLayer(nn.Module): def __init__(self): super().__init__() self.beta = nn.Parameter(torch.rand(1)) # 传播率 self.gamma = nn.Parameter(torch.rand(1)) # 恢复率 def forward(self, x, adj): # x: [S,I,R] S, I, R = x[...,0], x[...,1], x[...,2] dS = -self.beta * adj @ (S * I) dI = self.beta * adj @ (S * I) - self.gamma * I dR = self.gamma * I return torch.stack([dS, dI, dR], dim=-1)5.2 多源数据融合架构
整合四种关键数据源：
- 人口流动：手机信令数据
- 医疗资源：医院床位、诊所分布
- 环境因素：温度、湿度
- 防控措施：口罩佩戴率、社交限制
graph TD A[人口流动图] --> D[融合层] B[医疗资源图] --> D C[环境特征] --> D D --> E[STGNN编码器] E --> F[传播预测] E --> G[资源需求预测]实际案例表明，这种混合方法能将疫情峰值预测误差控制在±3天内，显著优于纯数据驱动方法。