ASTNN模型解析：如何用注意力时空神经网络解决道路级稀疏交通流预测难题

1次阅读

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

交通流预测是智能交通系统中的核心任务之一。传统的时空图卷积网络（STGCN）和扩散卷积循环网络（DCRNN）在处理道路级稀疏数据时，表现往往不尽如人意。特别是在检测器覆盖率低于 30% 的情况下，这些模型的预测精度会大幅下降。主要原因包括：

数据稀疏性 ：道路网络中许多路段缺乏有效的检测器数据，导致输入特征矩阵中存在大量零值。
长程依赖捕捉不足 ：传统模型难以有效捕捉稀疏数据中的长程时空依赖关系。
静态图结构限制 ：大多数模型依赖预定义的静态图结构，无法适应交通流的动态变化特性。

ASTNN 模型通过引入动态图构建和双重注意力机制，有效解决了上述问题。以下是 ASTNN 与几种主流模型的对比：

模型	动态图构建	注意力机制	稀疏数据处理能力
STGCN	❌	❌	一般
DCRNN	❌	❌	一般
时空 Transformer	❌	✔️	较好
ASTNN	✔️	✔️	优秀

ASTNN 通过动态图卷积层实时学习路段间的相关性，生成邻接矩阵。相较于 STGCN 和 DCRNN 的静态图结构，动态图能更好地适应交通流的时空变化。

ASTNN 包含时空双重注意力模块：

空间注意力 ：捕捉路段间的空间依赖关系。
时间注意力 ：建模交通流的时间动态特性。

这种双重注意力机制使得模型能够更有效地处理稀疏数据中的长程依赖。

动态图卷积层的核心是邻接矩阵的动态生成。以下是 PyTorch 实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class DynamicGraphConv(nn.Module):
    def __init__(self, in_features, out_features, nodes):
        super(DynamicGraphConv, self).__init__()
        self.W = nn.Parameter(torch.randn(in_features, out_features))
        self.nodes = nodes

    def forward(self, x):
        # x shape: [batch, nodes, features]
        batch_size = x.size(0)

        # 动态生成邻接矩阵
        adj = torch.bmm(x, x.transpose(1, 2))  # [batch, nodes, nodes]
        adj = F.softmax(adj, dim=-1)

        # 图卷积运算
        output = torch.bmm(adj, x)  # [batch, nodes, features]
        output = torch.matmul(output, self.W)  # [batch, nodes, out_features]

        return output

时空注意力模块通过矩阵运算实现：

class SpatioTemporalAttention(nn.Module):
    def __init__(self, hidden_dim):
        super(SpatioTemporalAttention, self).__init__()
        self.query = nn.Linear(hidden_dim, hidden_dim)
        self.key = nn.Linear(hidden_dim, hidden_dim)
        self.value = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        # x shape: [batch, nodes, timesteps, features]
        batch_size, nodes, timesteps, features = x.size()

        # 空间注意力
        x_reshaped = x.view(batch_size * nodes, timesteps, features)
        Q = self.query(x_reshaped)
        K = self.key(x_reshaped)
        V = self.value(x_reshaped)

        attn_weights = F.softmax(torch.bmm(Q, K.transpose(1, 2)) / (features ** 0.5), dim=-1)
        spatial_output = torch.bmm(attn_weights, V).view(batch_size, nodes, timesteps, features)

        # 时间注意力
        x_reshaped = x.permute(0, 2, 1, 3).contiguous().view(batch_size * timesteps, nodes, features)
        Q = self.query(x_reshaped)
        K = self.key(x_reshaped)
        V = self.value(x_reshaped)

        attn_weights = F.softmax(torch.bmm(Q, K.transpose(1, 2)) / (features ** 0.5), dim=-1)
        temporal_output = torch.bmm(attn_weights, V).view(batch_size, timesteps, nodes, features)
        temporal_output = temporal_output.permute(0, 2, 1, 3).contiguous()

        # 融合空间和时间注意力
        output = spatial_output + temporal_output
        return output

处理稀疏数据时，ASTNN 采用了以下技巧：

零值填充策略 ：对于缺失的路段数据，使用相邻时间步的均值填充。
mask 机制 ：在计算损失时，忽略缺失数据对应的输出。

def sparse_data_preprocess(data, mask):
    # data shape: [batch, nodes, timesteps, features]
    # mask shape: [batch, nodes, timesteps] (1 表示有数据，0 表示缺失)

    # 零值填充
    mean_values = data.sum(dim=0, keepdim=True) / (mask.sum(dim=0, keepdim=True) + 1e-6)
    filled_data = data * mask.unsqueeze(-1) + mean_values * (1 - mask.unsqueeze(-1))

    return filled_data