PyTorch实战：3种GNN模型搞定交通流量预测（附PEMS数据集处理技巧）

最新推荐文章于 2026-08-18 15:39:20 发布

原创

最新推荐文章于 2026-08-18 15:39:20 发布 · 570 阅读

从理论到实战：深度解析GNN在交通流量预测中的三种核心范式

如果你正在为城市交通的复杂性和不确定性感到头疼，那么图神经网络（GNN）可能是你一直在寻找的答案。交通流量预测远不止是简单的时间序列分析，它本质上是一个时空耦合的复杂系统问题。传感器节点（如交通检测器）之间的空间依赖关系，与每个节点自身随时间变化的动态模式，共同构成了预测的挑战。传统的RNN、LSTM在处理这类问题时，往往难以有效捕捉节点间的非欧几里得空间关系，而这正是GNN的天然优势。

过去几年，我参与过多个智慧城市交通项目，从最初的简单线性回归到后来的深度学习模型，踩过不少坑。一个深刻的体会是：模型的选择固然重要，但对数据的深刻理解和工程化实现的细节，往往才是项目成败的关键。本文将聚焦于三种最具代表性的GNN模型——GCN、ChebNet和GAT，并以经典的PEMS数据集为战场，带你从零构建一个完整的、可复现的交通流量预测项目。我们不仅会探讨模型原理，更会深入代码实现、数据预处理中的“魔鬼细节”，以及在实际部署中可能遇到的陷阱。

1. 理解战场：PEMS数据集深度剖析与预处理实战

在开始构建任何模型之前，我们必须先彻底了解我们的数据。PEMS（Performance Measurement System）数据集是交通预测领域的基准数据集，由加州交通部部署的感应环路检测器收集。它之所以经典，是因为它真实地反映了交通系统的复杂性：多节点、长时间跨度、多维度特征。

1.1 数据本质与结构洞察

以PEMS04为例，它包含了307个检测器连续59天的数据，采样间隔为5分钟。原始数据通常以.npz文件格式提供，加载后你会得到一个形状为(307, 16992, 3)的张量。这三个维度分别代表：

维度一（307）：空间维度，即检测器节点的数量。每个节点代表路网中的一个特定位置。
维度二（16992）：时间维度。59天 * 24小时/天 * 12个5分钟/小时 = 16992个时间步长。这是模型需要学习的时序模式所在。
维度三（3）：特征维度。通常包含流量（flow）、占有率（occupy） 和速度（speed）。

注意：在实际项目中，我们常常发现这三个特征之间存在高度的共线性。例如，流量和占有率通常呈正相关。因此，许多研究（包括本文的实践）会优先选择流量作为预测目标，因为它最直接地反映了道路的通行需求，也是上层应用（如信号控制、路径诱导）最关心的核心指标。

除了流量数据，我们还需要图结构信息，即邻接矩阵（Adjacency Matrix）。PEMS数据集通常提供一个distance.csv文件，包含from_node, to_node, distance三列。这个“距离”可以是实际的地理距离，也可以是行驶时间。我们的第一个关键决策就在这里：如何将物理距离转化为图卷积所需的邻接关系？

1.2 邻接矩阵构建：从物理连接到语义关联

构建邻接矩阵是GNN应用于交通预测的第一步，也是至关重要的一步。一个糟糕的图结构会直接导致模型无法学习有效的空间依赖。以下是几种常见的构建策略及其PyTorch实现：

策略一：阈值法 这是最直观的方法。我们设定一个距离阈值，如果两个节点间的距离小于该阈值，则认为它们相连。

import numpy as np
import torch

def build_adjacency_threshold(distance_df, node_ids, threshold=0.1):
    """
    基于距离阈值构建邻接矩阵。
    :param distance_df: DataFrame，包含'from','to','distance'列
    :param node_ids: 节点ID列表
    :param threshold: 归一化距离阈值，超过此值则不连接
    :return: 邻接矩阵 (N, N)
    """
    num_nodes = len(node_ids)
    adj = np.zeros((num_nodes, num_nodes), dtype=np.float32)
    node_to_idx = {node_id: idx for idx, node_id in enumerate(node_ids)}

    # 获取所有距离并归一化，以便设定通用阈值
    all_distances = distance_df['distance'].values
    max_dist = all_distances.max()
    min_dist = all_distances.min()
    normalized_distances = (all_distances - min_dist) / (max_dist - min_dist)

    for idx, row in distance_df.iterrows():
        i = node_to_idx[row['from']]
        j = node_to_idx[row['to']]
        dist_norm = normalized_distances[idx]
        if dist_norm <= threshold:
            adj[i, j] = 1.0
            adj[j, i] = 1.0  # 假设为无向图
    # 添加自连接
    np.fill_diagonal(adj, 1.0)
    return torch.from_numpy(adj)

策略二：K近邻法（KNN） 为每个节点选择距离最近的K个节点作为邻居。这种方法能保证每个节点都有固定数量的连接，避免了偏远节点的孤立问题。

from sklearn.neighbors import kneighbors_graph

def build_adjacency_knn(distance_matrix, k=5, mode='connectivity'):
    """
    使用K近邻法构建邻接矩阵。
    :param distance_matrix: 距离矩阵 (N, N)
    :param k: 近邻数量
    :param mode: 'connectivity'（0/1）或'distance'（保留距离）
    :return: 稀疏邻接矩阵 (N, N)
    """
    # 假设distance_matrix是一个对称的距离矩阵
    adj = kneighbors_graph(distance_matrix, n_neighbors=k, mode=mode, include_self=True)
    # 转换为稠密矩阵并确保对称（KNN图可能不对称）
    adj_dense = adj.toarray()
    adj_dense = np.maximum(adj_dense, adj_dense.T)  # 取并集确保对称
    return torch.from_numpy(adj_dense)

策略三：自适应学习法（高级） 这是目前前沿研究的方向。不依赖预定义的物理距离，而是让模型从数据中学习节点间的关联强度。这通常通过一个可学习的节点嵌入矩阵来实现，关联度由嵌入向量的相似度（如点积）决定。这种方法能捕捉到“功能相似性”（如两个商业区即使相距甚远，也可能有相似的交通模式）。

在我们的后续实验中，为了聚焦于模型对比，我们将采用最简单的阈值二值化法，并将所有有效连接的边权设为1。但请记住，在实际生产环境中，邻接矩阵的构建方式需要根据具体场景和数据特性进行仔细设计和验证。

1.3 数据标准化与序列构建

交通流量数据通常存在明显的日周期性和周周期性。为了帮助模型更好地学习，我们必须进行标准化。我推荐使用节点级（Node-wise）的Z-score标准化，即对每个传感器节点单独计算其历史数据的均值和标准差进行归一化。这能消除不同节点之间流量基数的差异。

def normalize_node_wise(data):
    """
    按节点进行Z-score标准化。
    :param data: 形状为 (N, T) 或 (N, T, F) 的流量数据
    :return: 标准化后的数据，以及每个节点的均值和标准差（用于逆变换）
    """
    if data.ndim == 2:
        means = data.mean(axis=1, keepdims=True)
        stds = data.std(axis=1, keepdims=True)
        stds[stds == 0] = 1.0  # 防止除零
        norm_data = (data - means) / stds
        return norm_data, means, stds
    elif data.ndim == 3:
        # 对于多特征，通常对每个特征单独标准化
        N, T, F = data.shape
        norm_data = np.zeros_like(data)
        means = np.zeros((N, 1, F))
        stds = np.zeros((N, 1, F))
        for f in range(F):
            feat_data = data[:, :, f]
            mean_f = feat_data.mean(axis=1, keepdims=True)
            std_f = feat_data.std(axis=1, keepdims=True)
            std_f[std_f == 0] = 1.0
            norm_data[:, :, f] = (feat_data - mean_f) / std_f
            means[:, :, f] = mean_f
            stds[:, :, f] = std_f
        return norm_data, means, stds

接下来，我们需要将长时间序列切割成模型可处理的样本。假设我们使用过去12个时间步（即1小时）的数据来预测未来3个时间步

标签