如何用Transformer网络提升交通流量预测准确率？实战PeMS-Bay数据集

从理论到实战：用Transformer重塑交通流量预测的工程实践

如果你正在为城市交通预测的准确性而头疼，觉得传统模型在早晚高峰的剧烈波动面前总是“慢半拍”，或者面对海量传感器数据时，模型难以捕捉那些瞬息万变的时空关联，那么这篇文章或许能为你打开一扇新窗。交通流量预测，这个智能交通系统的核心难题，正随着Transformer架构的崛起，迎来一场深刻的变革。我们不再仅仅依赖循环神经网络（RNN）或卷积神经网络（CNN）的局部感知，而是转向一种能同时、动态地理解全局时空依赖关系的强大工具。

本文将聚焦于一个经典的实战场景：基于PeMS-Bay数据集的交通流量预测。我们将绕过繁琐的理论推导，直接从工程落地角度切入，分享如何将Transformer，特别是其时空变体，应用于这个真实世界的数据集。你会看到从原始数据清洗、模型架构选型、关键参数调优，到效果评估与对比的完整链条。我们的目标很明确：提升预测准确率，尤其是应对长期、复杂多变的交通模式。无论你是智能交通系统的开发者，还是希望将前沿AI模型应用于时空预测的数据科学家，这里都有可供直接参考的实践经验。

1. 理解挑战：为什么传统方法在交通预测上力不从心？

在深入代码之前，我们必须先理解问题的本质。交通流量数据本质上是一种时空图数据：每个交通传感器是一个节点，道路连接构成了边，而每个节点在不同时间点采集的速度、流量等信息，构成了随时间演变的图信号。预测未来流量，就是基于历史的图信号序列，推断未来的图信号。

传统方法，如ARIMA、卡尔曼滤波等时间序列模型，擅长捕捉单一节点的时序规律，但完全忽视了空间维度——下游路口的事故显然会影响上游的流速，而这是ARIMA无法建模的。后来的方法引入了图神经网络（GNN），例如STGCN、DCRNN，它们通过图卷积来捕捉空间依赖，再结合时序模型（如TCN、GRU）处理时间维度。这固然是一大进步，但仍存在两个核心瓶颈：

1. 静态的空间依赖假设：大多数GNN-based模型使用一个预定义的、固定的邻接矩阵（通常基于传感器间的物理距离或道路连接）。然而，实际的空间依赖是动态且具有方向性的。早高峰时，进城方向的关联性远强于出城方向；一场突发事故的影响范围会随时间扩散和收缩。固定图结构无法刻画这种时变特性。
2. 有限的时间感受野：RNN类模型存在梯度消失/爆炸问题，难以学习长序列依赖；基于卷积的时序模型（如TCN）需要堆叠很多层才能获得较大的感受野，计算效率低，且难以自适应不同长度的输入序列。

这就引出了Transformer的核心优势：其自注意力（Self-Attention）机制能够直接计算序列中任意两个元素之间的关联权重，无论它们相距多远。将其适配到时空图数据上，我们就能让模型动态地学习节点与节点之间（空间）、时间步与时间步之间（时间）的依赖关系，且这种学习是数据驱动的，无需先验的固定假设。

提示：理解“动态空间依赖”是掌握时空Transformer价值的关键。可以想象，自注意力机制会为每一对传感器节点在每一个时间步都计算一个“关联度分数”，这个分数完全由当前时刻的交通状态数据决定，从而实时反映路网中不断变化的相互影响。

2. 实战起点：PeMS-Bay数据集深度处理与特征工程

任何优秀的模型都建立在高质量的数据基础之上。PeMS-Bay是交通预测领域广泛使用的基准数据集，包含了加州湾区325个传感器站6个月（2017年1月至5月）的交通速度数据，采样间隔为5分钟。

2.1 数据获取与初步探索

首先，你需要从官方渠道获取数据。通常包含两个核心文件：pems-bay.h5（交通速度数据）和adj_mx_bay.pkl（基于距离计算的邻接矩阵）。

import pandas as pd
import numpy as np
import h5py
import pickle

# 加载速度数据
with h5py.File('pems-bay.h5', 'r') as f:
    speed_data = f['speed'][:]  # 形状: (时间步数, 节点数)
print(f"数据形状: {speed_data.shape}")  # 例如 (52116, 325)

# 加载邻接矩阵（基于距离的高斯核权重）
with open('adj_mx_bay.pkl', 'rb') as f:
    adj_matrix = pickle.load(f)

初步探索时，务必检查缺失值、异常值（如速度为0或极大值）以及数据的整体分布。绘制几个关键节点的时序曲线，观察其日周期、周周期特性。

2.2 构建时空图数据结构

我们需要将数据组织成模型所需的格式。一个常见的做法是构建一个三维张量 X，其维度为 (样本数, 历史时间步长, 节点数, 特征数)。这里特征数初始为1（速度）。

def create_dataset(data, seq_len, pred_len):
    """
    将一维时序数据转换为监督学习样本。
    data: 形状 (总时间步, 节点数)
    seq_len: 历史窗口长度
    pred_len: 预测未来长度
    """
    num_samples, num_nodes = data.shape
    samples = []
    labels = []
    for i in range(num_samples - seq_len - pred_len + 1):
        samples.append(data[i:i+seq_len, :])  # 历史窗口
        labels.append(data[i+seq_len:i+seq_len+pred_len, :])  # 未来窗口
    return np.array(samples), np.array(labels)

# 示例：使用12个历史步（60分钟）预测未来3步（15分钟）
seq_length = 12
pred_length = 3
X, y = create_dataset(speed_data.T, seq_length, pred_length) # 注意转置，使节点数在第二维
print(f"样本集X形状: {X.shape}")  # (N, 12, 325, 1)
print(f"标签集y形状: {y.shape}")  # (N, 3, 325, 1)

2.3 关键预处理技巧与陷阱规避

1. 标准化（Normalization）: 对每个传感器节点的速度序列分别进行Z-Score标准化。这是至关重要的一步，能加速模型收敛并提升性能。切记要使用训练集的均值和标准差来转换验证集和测试集，避免数据泄露。

   from sklearn.preprocessing import StandardScaler
   scalers = {}
   data_normalized = np.zeros_like(speed_data.T)
   for node in range(num_nodes):
       scaler = StandardScaler()
       # 假设已划分好train_idx
       scaler.fit(speed_data.T[train_idx, node].reshape(-1, 1))
       data_normalized[:, node] = scaler.transform(speed_data.T[:, node].reshape(-1, 1)).flatten()
       scalers[node] = scaler
   

2. 图结构的处理: 原始邻接矩阵是基于距离的高斯核，权重是连续的。对于某些模型，可能需要将其二值化（例如，设定一个阈值，大于阈值的边保留为1，否则为0）或进行归一化（如对称归一化拉普拉斯矩阵）。对于时空Transformer，我们通常保留原始的连续权重矩阵，因为它能提供更丰富的空间先验信息。

3. 时间特征嵌入: 除了速度，周期性的时间信息是强大的特征。可以编码小时、星期几、是否为节假日等。

   def create_time_features(timestamps):
       """为每个时间步创建时间特征"""
       features = []
       for ts in timestamps:
           dt = pd.to_datetime(ts, unit='s')  # 假设时间戳为秒
           features.append([
               dt.hour / 23.0,  # 归一化小时
               dt.dayofweek / 6.0, # 归一化星期几
               float(dt.month in [7,8]), # 是否为暑假（示例）
           ])
       return np.array(features)
   

   这些时间特征可以作为额外的通道（特征数>1）与速度数据拼接，输入模型。

4. 数据集划分: 务必按时间顺序划分训练集、验证集和测试集（例如，前70%训练，中间15%验证，最后15%测试），绝不能随机打乱，以模拟真实的时间序列预测场景。

3. 构建时空Transformer模型：核心模块拆解

现在进入核心环节。我们将构建一个简化但功能完整的时空Transformer网络（STTN）。其核心思想是：空间Transformer模块负责在每一个时间步内部，动态计算所有节点间的注意力权重；时间Transformer模块负责在每个节点维度上，跨越历史时间步计算注意力权重。两者交替或并行堆叠。

3.1 空间Transformer模块：学习动态的路网关联

空间Transformer的目标是替换掉传统的、固定的图卷积。对于某个时间步 t 的所有节点特征 X_t ∈ R^(N×d)，我们通过自注意力计算一个动态的邻接矩阵 A_dynamic_t。

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

class SpatialTransformerLayer(nn.Module):
    def __init__(self, node_num, in_features, out_features, num_heads=4, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.node_num = node_num
        self.in_features = in_features
        self.out_features = out_features
        self.head_dim = out_features // num_heads

        assert out_features % num_heads == 0, "out_features must be divisible by num_heads"

        # 线性变换得到Q, K, V
        self.W_q = nn.Linear(in_features, out_features)
        self.W_k = nn.Linear(in_features, out_features)
        self.W_v = nn.Linear(in_features, out_features)
        self.fc_out = nn.Linear(out_features, out_features)
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(out_features)

        # 可选的静态空间先验注入（如通过预定义的邻接矩阵）
        self.static_adj = None  # 可以初始化为一个可学习的参数或固定的矩阵

    def forward(self, x, static_adj=None):
        """
        x: 输入张量，形状为 (batch_size, node_num, in_features)
        static_adj: 可选的静态邻接矩阵先验，形状 (node_num, node_num)
        返回: 输出张量，形状 (batch_size, node_num, out_features)
        """
        batch_size = x.shape[0]
        # 1. 线性投影得到Q, K, V
        Q = self.W_q(x).view(batch_size, self.node_num, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.W_k(x).view(batch_size, self.node_num, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.W_v(x).view(batch_size, self.node_num, self.num_heads, self.head_dim).transpose(1, 2)

        # 2. 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)  # (B, H, N, N)

        # 3. （关键步骤）融入静态空间先验知识
        if static_adj is not None:
            # 将静态邻接矩阵扩展到多头和批次维度，并与注意力分数结合
            # 例如，可以用加性掩码或乘性权重
            static_adj_expanded = static_adj.unsqueeze(0).unsqueeze(0)  # (1,1,N,N)
            scores = scores + static_adj_expanded  # 简单加性融合，可替换为更复杂的方式

        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # 4. 应用注意力权重到V
        context = torch.matmul(attention_weights, V)  # (B, H, N, head_dim)
        context = context.transpose(1, 2).contiguous().view(batch_size, self.node_num, self.out_features)

        # 5. 输出投影与残差连接
        output = self.fc_out(context)
        output = self.dropout(output)
        output = self.layer_norm(output + x)  # 残差连接
        return output, attention_weights.detach()  # 返回注意力权重用于分析

这个模块的工程要点：

• 多头注意力：允许模型在不同的表示子空间中共同关注来自不同位置的信息，这对于捕捉复杂的、多因素驱动的空间依赖（如上下游关系、车流相似性、道路等级等）至关重要。
• 静态先验融合：完全依赖数据学习动态关联可能在小数据集上不稳定。将预定义的基于距离的邻接矩阵作为先验知识（例如，以加性偏置或乘性掩码的形式）注入注意力计算，可以引导模型学习，提升训练效率和泛化能力。
• 计算效率：节点数 N 较大时（如PeMS-Bay的325个），计算 N×N 的注意力矩阵开销很大。实践中可以考虑节点采样、局部注意力窗口（只计算每个节点k近邻的注意力）或线性注意力等优化技巧。

3.2 时间Transformer模块：捕捉长程时序模式

时间Transformer的结构与标准Transformer编码器层类似，但作用在时间维度上。对于某个节点 n 的历史特征序列 X_n ∈ R^(T×d)，我们计算时间步之间的注意力。

class TemporalTransformerLayer(nn.Module):
    def __init__(self, seq_len, in_features, out_features, num_heads=4, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.seq_len = seq_len
        self.in_features = in_features
        self.out_features = out_features
        self.head_dim = out_features // num_heads

        # 位置编码：为时间步添加顺序信息
        self.pos_encoder = PositionalEncoding(out_features, dropout, max_len=5000)

        # 自注意力层
        self.self_attn = nn.MultiheadAttention(out_features, num_heads, dropout=dropout, batch_first=True)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(out_features)

        # 前馈网络
        self.ffn = nn.Sequential(
            nn.Linear(out_features, out_features * 4),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(out_features * 4, out_features),
        )
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(out_features)

    def forward(self, src):
        """
        src: 输入张量，形状 (batch_size, seq_len, node_num, in_features)
        为简化，我们先处理单个节点维度，实际需遍历节点或重塑张量。
        这里假设输入已重塑为 (batch_size * node_num, seq_len, in_features)
        """
        # 添加位置编码
        src = self.pos_encoder(src)

        # 自注意力
        src2, attn_weights = self.self_attn(src, src, src)
        src = src + self.dropout1(src2)
        src = self.norm1(src)

        # 前馈网络
        src2 = self.ffn(src)
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src, attn_weights.detach()

class PositionalEncoding(nn.Module):
    """标准正弦/余弦位置编码"""
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)  # (max_len, 1, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

时间模块的注意点：

• 因果掩码（Causal Masking）: 对于预测任务，在训练时，时间Transformer的注意力层通常需要应用因果掩码，确保在预测时间步 t 时，只能看到 t 之前（含 t）的信息，防止信息泄露。但在时空预测中，我们通常使用历史窗口的所有时间步来预测未来窗口，历史窗口内部是已知的，因此编码器部分可以不使用因果掩码，允许所有历史时间步相互关注，以更好地提取特征。
• 位置编码的必要性：自注意力机制本身不具备序列顺序信息，必须通过位置编码注入时间步的先后关系。

3.3 模型整合与预测头

将空间和时间模块组合起来，形成完整的STTN。一种常见的架构是“时空块”的堆叠：每个块内，先进行空间注意力聚合，再进行时间注意力聚合，并加入残差连接。

class SpatioTemporalBlock(nn.Module):
    def __init__(self, node_num, seq_len, in_features, out_features, spatial_heads, temporal_heads, dropout):
        super().__init__()
        self.spatial_transformer = SpatialTransformerLayer(node_num, in_features, out_features, spatial_heads, dropout)
        self.temporal_transformer = TemporalTransformerLayer(seq_len, out_features, out_features, temporal_heads, dropout)
        # 可能还需要一个前馈网络或额外的层归一化

    def forward(self, x, static_adj):
        # x shape: (B, T, N, C)
        B, T, N, C = x.shape
        # 空间变换：在每一个时间步内部进行
        spatial_out = []
        spatial_attns = []
        for t in range(T):
            s_out, s_attn = self.spatial_transformer(x[:, t, :, :], static_adj)
            spatial_out.append(s_out)
            spatial_attns.append(s_attn)
        spatial_out = torch.stack(spatial_out, dim=1)  # (B, T, N, C)

        # 时间变换：对每一个节点序列进行
        # 重塑: (B, T, N, C) -> (B*N, T, C)
        temporal_in = spatial_out.permute(0, 2, 1, 3).contiguous().view(B*N, T, -1)
        temporal_out, temporal_attn = self.temporal_transformer(temporal_in)
        # 重塑回来: (B*N, T, C) -> (B, T, N, C)
        temporal_out = temporal_out.view(B, N, T, -1).permute(0, 2, 1, 3).contiguous()

        # 残差连接
        output = temporal_out + x
        return output, spatial_attns, temporal_attn

class STTN(nn.Module):
    def __init__(self, node_num, seq_len, pred_len, in_dim=1, hidden_dim=64, num_blocks=2, spatial_heads=4, temporal_heads=4, dropout=0.1):
        super().__init__()
        self.node_num = node_num
        self.seq_len = seq_len
        self.pred_len = pred_len

        # 输入嵌入层（可包含时间特征）
        self.input_proj = nn.Linear(in_dim, hidden_dim)

        # 堆叠时空块
        self.st_blocks = nn.ModuleList()
        for _ in range(num_blocks):
            self.st_blocks.append(SpatioTemporalBlock(node_num, seq_len, hidden_dim, hidden_dim, spatial_heads, temporal_heads, dropout))

        # 预测头：简单的全连接层或卷积层
        self.output_layer = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim * 2, pred_len)  # 直接输出未来pred_len个时间步
        )

    def forward(self, x, static_adj):
        # x: (B, T, N, C_in)
        x = self.input_proj(x)  # (B, T, N, hidden_dim)

        spatial_attn_list = []
        temporal_attn_list = []
        for block in self.st_blocks:
            x, s_attn, t_attn = block(x, static_adj)
            spatial_attn_list.append(s_attn)
            temporal_attn_list.append(t_attn)

        # 取最后一个时间步的特征用于预测（常见做法）
        last_step_features = x[:, -1, :, :]  # (B, N, hidden_dim)
        # 或者使用所有时间步特征的平均/最大池化

        # 预测每个节点未来pred_len个时间步的值
        prediction = self.output_layer(last_step_features)  # (B, N, pred_len)
        prediction = prediction.transpose(1, 2).unsqueeze(-1)  # (B, pred_len, N, 1)

        return prediction, spatial_attn_list, temporal_attn_list

4. 训练、调参与效果分析：通往高精度的路径

模型搭建完毕，真正的挑战才刚刚开始。如何训练好这个参数众多的网络，并使其在PeMS-Bay上达到最优性能？

4.1 损失函数与评估指标

交通预测通常使用平均绝对误差（MAE） 或均方误差（MSE） 作为损失函数。MAE对异常值不那么敏感，而MSE会给予大误差更大的惩罚，可能导致模型更倾向于平滑预测。

criterion = nn.L1Loss()  # MAE
# 或 criterion = nn.MSELoss()

# 评估时，除了MAE，还应计算以下指标：
def compute_metrics(preds, targets):
    mae = torch.abs(preds - targets).mean()
    mse = ((preds - targets) ** 2).mean()
    rmse = torch.sqrt(mse)
    # MAPE (Mean Absolute Percentage Error) 需注意分母为零的情况
    mape = torch.abs((preds - targets) / (targets + 1e-5)).mean() * 100
    return {'MAE': mae.item(), 'RMSE': rmse.item(), 'MAPE': mape.item()}

4.2 超参数调优经验

时空Transformer包含大量超参数，系统性的调优至关重要。以下是一个基于经验的调参优先级和范围参考：

超参数	建议搜索范围/策略	影响说明
学习率 (Learning Rate)	[1e-4, 5e-4, 1e-3]	Transformer模型通常需要较小的学习率或带热启动（Warmup）的策略。使用AdamW优化器。
隐藏层维度 (hidden_dim)	[32, 64, 128]	影响模型容量。PeMS-Bay上64是一个不错的起点。太大易过拟合，太小欠拟合。
时空块数量 (num_blocks)	[1, 2, 3]	增加深度可以增强模型表达能力，但也会增加训练难度和过拟合风险。2-3层通常足够。
空间注意力头数 (spatial_heads)	[2, 4, 8]	多头允许模型关注不同的空间依赖模式（如距离、流向、功能相似性）。4或8是常见选择。
时间注意力头数 (temporal_heads)	[2, 4]	时间模式相对空间可能简单，头数可以少一些。
历史序列长度 (seq_len)	[12, 24, 36, 48]	对应1, 2, 3, 4小时。更长的序列提供更多上下文，但计算成本增加。需要平衡。
Dropout 比率	[0.1, 0.2, 0.3]	防止过拟合的关键正则化手段。在模型较大或数据较少时提高比率。
批处理大小 (Batch Size)	[16, 32, 64]	受GPU内存限制。在内存允许下，较大的Batch Size通常使训练更稳定。

调参实战建议：

1. 从基准开始：先固定一组保守参数（如lr=1e-3, hidden_dim=64, blocks=2, heads=4）训练一个基础模型，确保代码能跑通且损失下降。
2. 学习率与优化器优先：使用学习率查找器（LR Finder）或余弦退火等自适应调度策略。AdamW搭配CosineAnnealingLR在Transformer类模型上表现通常很好。
3. 宽度 vs 深度：如果模型欠拟合（训练集误差也高），尝试增加hidden_dim或num_blocks。如果过拟合（训练集误差低，验证集误差高），首先增加dropout，或减少hidden_dim，添加更多的正则化（如权重衰减）。
4. 序列长度实验：这是一个对性能影响巨大的参数。尝试不同的seq_len，观察验证集指标的变化。有时更长的历史未必带来更好的预测，因为模型可能被无关的早期信息干扰。
5. 利用注意力可视化：训练后，绘制学到的空间和时间注意力权重图。这不仅能验证模型是否学到了有意义的模式（例如，空间注意力是否聚焦于物理相邻或流量相似的节点，时间注意力是否关注了周期性的早高峰时段），还能帮助诊断问题（例如，注意力是否过于分散或集中在无关的节点/时间上）。

4.3 与基线模型的对比实验

为了客观评估STTN的效果，必须在相同的训练/验证/测试集划分、相同的预处理流程下，与主流基线模型进行对比。常见的基线包括：

• 历史平均（HA）
• ARIMA
• FC-LSTM: 全连接LSTM，忽略空间关系。
• STGCN：结合图卷积和时序卷积的经典时空模型。
• DCRNN：使用扩散图卷积和循环门控单元的模型。
• Graph WaveNet：结合自适应邻接矩阵和扩张因果卷积的强基线。

在PeMS-Bay数据集上，一个训练良好的STTN模型，在预测未来1小时（12个时间步）的任务上，其MAE指标通常能显著优于FC-LSTM和STGCN，并与Graph WaveNet这样的先进模型达到可比甚至更优的水平，尤其是在预测时长增加时，其优势可能更明显，这得益于其对长程依赖的建模能力。

4.4 工程化部署的考量

当模型在离线评估中表现良好后，需要考虑线上部署：

• 推理速度：Transformer的自注意力计算复杂度与序列长度的平方成正比。对于实时预测，如果历史窗口很长，需要考虑优化，如使用线性注意力、局部注意力或模型蒸馏。
• 模型轻量化：通过剪枝、量化等技术减小模型体积，以适应边缘设备部署。
• 持续学习：交通模式会随时间缓慢变化（新路开通、政策调整）。需要设计机制，使模型能够在不遗忘旧知识的情况下，利用新数据进行增量更新。

在我自己的几个交通预测项目中，踩过最大的一个坑是忽视了数据泄露——在全局做标准化时错误地使用了全部数据的均值和方差。这导致验证集和测试集指标虚高，上线后效果大打折扣。另一个经验是，空间注意力的可视化不仅是一个分析工具，更是说服领域专家和项目决策者的有力武器。当你展示出模型自动识别出了某个关键瓶颈路口对其上游一片区域的强影响力时，其可信度会大大增加。

Transformer为交通流量预测带来了新的可能性，它通过动态的、全局的注意力机制，更细腻地刻画了城市路网中复杂的时空交互。从PeMS-Bay这个“练兵场”出发，理解数据、精心构建模型、耐心调参、严谨评估，这套方法论可以迁移到更广泛的时空序列预测问题中，例如电网负荷预测、人群流量预测、环境指标监测等。技术的价值在于解决实际问题，而Transformer正为我们提供了一把更精准的钥匙。