保姆级教程：用Python处理METR-la交通数据集，搞定12步预测的输入输出格式
从零解析METR-la交通数据集：12步预测的完整数据工程指南
当你第一次打开METR-la交通数据集时，那个看似简单的N×T矩阵背后隐藏着多少未解之谜？作为时空预测领域的经典基准数据集，METR-la的原始.h5文件就像一座未经雕琢的矿山，而本文将带你用Python这把精密工具，完成从原始数据到模型可消化格式的完整蜕变过程。
1. 理解METR-la数据集的本质
METR-la数据集记录了洛杉矶高速公路207个传感器在4个月内的交通速度数据，原始格式为34272个时间点×207个传感器的二维矩阵。但现代时空图神经网络(ST-GNN)需要的是序列到序列的输入输出范式——用过去12个时间步预测未来12个时间步。这就产生了三个核心问题：
- 如何从N×T矩阵中提取出N×12的序列片段？
- 如何处理时间边界处的数据截断？
- 如何合理分割训练集、验证集和测试集？
关键洞察：交通数据具有明显的时空双重特性——空间上传感器构成图结构，时间上速度变化呈现周期性。预处理必须保留这两种特性。
让我们先建立基础环境：
import numpy as np import pandas as pd import os from datetime import datetime # 确保可复现性 np.random.seed(42)2. 数据加载与初步探索
原始数据存储在HDF5格式中，Pandas提供了便捷的读取方式：
def load_metr_la_data(file_path='metr-la.h5'): df = pd.read_hdf(file_path) print(f"数据集形状：{df.shape}") # 应输出(34272, 207) print(f"时间范围：{df.index.min()} 至 {df.index.max()}") return df执行后会看到类似输出：
数据集形状：(34272, 207) 时间范围：2012-03-01 00:00:00 至 2012-06-30 23:55:00数据特性分析：
- 时间分辨率为5分钟（288个时间点/天）
- 包含3%的缺失值（用NaN表示）
- 速度单位是mph(英里/小时)
3. 构建时间序列偏移量体系
时空预测的核心在于构建时间滑动窗口。我们需要定义两组关键偏移量：
# 输入序列：过去12个时间步（含当前时刻） x_offsets = np.arange(-11, 1, 1) # [-11, -10,..., 0] # 预测目标：未来12个时间步（不含当前时刻） y_offsets = np.arange(1, 13, 1) # [1, 2,..., 12]这种设计意味着：
- 每个样本x[i]包含t-11到t时刻的数据
- 对应标签y[i]包含t+1到t+12时刻的数据
技术细节：偏移量选择直接影响模型学习的时间依赖范围。12步对应1小时跨度，适合捕捉交通流的短时演变规律。
4. 时间特征工程实战
除了交通速度，时间本身也是重要特征。我们需要提取两种时间信息：
def create_time_features(df): # 一天中的相对位置（0-1之间） time_ind = (df.index.values - df.index.values.astype("datetime64[D]")) time_ind = time_ind / np.timedelta64(1, "D") # 星期几（0-6） day_of_week = df.index.dayofweek.values / 7 return np.stack([time_ind, day_of_week], axis=-1)特征拼接后的数据维度变化：
原始数据：(34272, 207, 1) → 添加时间后：(34272, 207, 3)5. 滑动窗口生成算法
这是最关键的步骤，将长序列切割为训练样本：
def generate_sequence_samples(data, x_offsets, y_offsets): num_samples, num_nodes, _ = data.shape min_t = abs(min(x_offsets)) # 11 max_t = num_samples - max(y_offsets) # 34272-12=34260 x, y = [], [] for t in range(min_t, max_t): x.append(data[t + x_offsets]) # 12×207×3 y.append(data[t + y_offsets]) # 12×207×3 return np.stack(x), np.stack(y)内存优化技巧： 对于大型数据集，建议使用生成器而非一次性加载所有样本：
def sequence_generator(data, batch_size=64): # 实现批处理生成逻辑 while True: for i in range(0, len(data), batch_size): yield data[i:i+batch_size]6. 数据集分割策略
交通数据具有强时间相关性，必须按时间顺序分割：
| 数据集 | 比例 | 样本数 | 时间范围 |
|---|---|---|---|
| 训练集 | 70% | 23,974 | 前3个月 |
| 验证集 | 10% | 3,425 | 第4个月前三周 |
| 测试集 | 20% | 6,850 | 最后一周 |
实现代码：
def split_dataset(x, y, train_ratio=0.7, val_ratio=0.1): num_samples = x.shape[0] num_train = int(num_samples * train_ratio) num_val = int(num_samples * val_ratio) return { 'train': (x[:num_train], y[:num_train]), 'val': (x[num_train:num_train+num_val], y[num_train:num_train+num_val]), 'test': (x[num_train+num_val:], y[num_train+num_val:]) }7. 数据标准化与缺失值处理
交通数据需要两种标准化方法：
全局标准化（适合静态模型）：
scaler = StandardScaler() data_scaled = scaler.fit_transform(df.values.reshape(-1, 1)).reshape(df.shape)局部标准化（适合自适应模型）：
def rolling_normalize(data, window_size=288): # 1天窗口 rolling_mean = data.rolling(window=window_size, min_periods=1).mean() rolling_std = data.rolling(window=window_size, min_periods=1).std() return (data - rolling_mean) / (rolling_std + 1e-8)
对于缺失值，推荐采用时空混合插补：
from sklearn.impute import KNNImputer # 空间维度插补（利用相邻传感器） spatial_imputer = KNNImputer(n_neighbors=5) data_spatial_filled = spatial_imputer.fit_transform(df.values) # 时间维度插补（利用历史同期数据） temporal_imputer = KNNImputer(n_neighbors=3) data_final = temporal_imputer.fit_transform(data_spatial_filled.T).T8. 高效存储方案设计
NPZ格式比HDF5更轻量且兼容性好：
def save_as_npz(data_dict, output_dir): os.makedirs(output_dir, exist_ok=True) for name, (x, y) in data_dict.items(): np.savez_compressed( os.path.join(output_dir, f"{name}.npz"), x=x.astype(np.float32), # 节省存储空间 y=y.astype(np.float32), x_offsets=x_offsets, y_offsets=y_offsets )文件结构示例：
output_dir/ ├── train.npz ├── val.npz └── test.npz9. 数据可视化验证
在预处理完成后，必须进行可视化校验：
import matplotlib.pyplot as plt def plot_sample(x, y, node_idx=0, feature_idx=0): plt.figure(figsize=(12, 6)) plt.plot(x[:, node_idx, feature_idx], label='Input') plt.plot(range(12, 24), y[:, node_idx, feature_idx], label='Target') plt.legend() plt.show()典型检查点：
- 输入输出序列的时间对齐是否正确
- 标准化后的数据是否保持原有模式
- 缺失值插补是否自然
10. 与主流框架的集成
处理后的数据可直接用于流行库：
PyTorch数据加载：
from torch.utils.data import Dataset class TrafficDataset(Dataset): def __init__(self, npz_file): data = np.load(npz_file) self.x = data['x'] self.y = data['y'] def __len__(self): return len(self.x) def __getitem__(self, idx): return torch.FloatTensor(self.x[idx]), torch.FloatTensor(self.y[idx])DGL图神经网络输入：
import dgl def build_spatial_graph(num_nodes=207): # 基于距离构建图结构 src = [...] # 源节点列表 dst = [...] # 目标节点列表 return dgl.graph((src, dst), num_nodes=num_nodes)11. 性能优化技巧
处理大规模交通数据时，这些技巧可提升效率：
内存映射技术：
x = np.load('train.npz', mmap_mode='r')['x']并行化处理：
from joblib import Parallel, delayed def process_chunk(chunk): return generate_sequence_samples(chunk, x_offsets, y_offsets) results = Parallel(n_jobs=4)(delayed(process_chunk)(chunk) for chunk in np.array_split(data, 4))数据压缩策略对比：
| 格式 | 读取速度 | 存储大小 | 兼容性 |
|---|---|---|---|
| NPZ | ★★★★ | ★★★ | ★★★★★ |
| HDF5 | ★★★ | ★★★★ | ★★★★ |
| Feather | ★★★★★ | ★★ | ★★★ |
12. 常见问题解决方案
问题1：序列长度不匹配
- 检查
max_t计算是否准确：max_t = num_samples - max(y_offsets)
问题2：内存不足
- 使用分块处理：
for chunk in pd.read_hdf('metr-la.h5', chunksize=10000)
问题3：预测结果滞后
- 验证偏移量方向：
y_offsets应为正数序列
问题4：模型无法收敛
- 检查标准化是否反向应用了：测试集必须使用训练集的均值和方差
在真实项目中，我发现最易出错的是时间偏移量的方向设置。曾经因为将y_offsets误设为np.arange(0, 12)导致模型学习到恒等映射而非预测能力。另一个经验是，对于长期预测（如预测未来1小时以上），建议在时间特征中加入星期几和节假日标志，这对捕捉交通模式的宏观周期非常有效。