ASTNN实战：基于注意力时空神经网络的稀疏交通流预测解决方案

1次阅读

没有评论

共计 1877 个字符，预计需要花费 5 分钟才能阅读完成。

交通流预测在城市规划、网约车调度等领域很常见，但实际操作中总会遇到两个头疼的问题：

数据稀疏性 ：很多道路的传感器数据时有时无，像郊区道路可能几个小时才有一辆车经过
时空异质性 ：早高峰的市区和晚间的工业区流量模式完全不同，但传统模型难以捕捉这种动态变化

我试过用 ARIMA 做预测，结果发现：

只能处理单一路径的时间序列，完全忽略路网拓扑
遇到缺失数据要么插值要么丢弃，预测误差能到 30% 以上

后来用 ST-ResNet 这类深度学习模型，虽然效果好些，但还是有问题：

卷积核大小固定，无法自适应不同区域的路网密度
对所有时间片同等对待，没法突出早晚高峰的特殊性

ASTNN 模型的核心思路是模拟人类司机对路网的认知方式：

空间层面 ：通过邻居道路推断当前道路状态（就像司机知道前方堵车会绕道）
时间层面 ：重点记忆周期性规律（比如每周五晚高峰必堵）
注意力机制 ：动态分配权重（事故路段要特别关注）

这里用到了改进的邻接矩阵，不是简单的 0 / 1 表示是否相邻，而是包含：

道路等级差异（高速路与辅路权重不同）
实际车流转移概率（从 A 路到 B 路的历史转向比例）

# PyTorch 图卷积实现示例
class GraphConv(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)

    def forward(self, x, adj):
        # x: [batch, nodes, features]
        # adj: [nodes, nodes] 动态调整的邻接矩阵
        return F.relu(torch.bmm(adj, self.linear(x)))

类似 LSTM 但做了两点优化：

输入门增加稀疏性掩码（对缺失数据自动降权）
遗忘门与周 / 日周期联动（周一早 8 点的记忆优先保留）

这里最精妙的是空间和时间注意力的交叉计算：

# 注意力头计算示例
def scaled_dot_product_attention(Q, K, V):
    """
    Q: 查询向量 [batch, heads, nodes, dim]
    K: 键向量   [batch, heads, time_steps, dim]
    V: 值向量   [batch, heads, time_steps, dim]
    """
    attn = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(K.size(-1))
    attn = F.softmax(attn, dim=-1)
    return torch.matmul(attn, V)

动态邻接矩阵 ：每小时根据实际流量重新计算道路关联度
缺失值掩码 ：对无数据时段，用相邻路段加权填充并打上标记
课程学习策略 ：训练时先学习数据密集区域，逐步加入稀疏样本

原始 GPS 数据按 5 分钟间隔划分
将城市网格化为 500m×500m 区域
对每个网格：
统计车辆数作为流量值
计算进出相邻网格的转移概率

# 关键训练参数
batch_size = 32  # 太大容易 OOM
window_size = 12  # 1 小时历史数据
heads = 4         # 注意力头数

# 使用梯度检查点节省显存
from torch.utils.checkpoint import checkpoint
model = checkpoint(model, inputs)