当前位置：首页 > article >正文

介绍如何基于现有的可运行STGCN（Spatial-Temporal Graph Convolutional Network）模型代码进行交通流预测的改动

article 2025/3/10 17:10:58

下面将详细介绍如何基于现有的可运行STGCN（Spatial-Temporal Graph Convolutional Network）模型代码进行交通流预测的改动。STGCN是一种用于处理时空数据的深度学习模型，非常适合交通流预测任务。

步骤概述

数据准备：加载和预处理交通流数据，构建图结构。
模型修改：确保STGCN模型适用于交通流预测任务。
训练和评估：训练模型并评估其在交通流预测上的性能。

详细步骤

1. 数据准备

首先，你需要准备交通流数据，通常以时间序列的形式表示，同时构建图结构来表示交通网络中节点之间的关系。以下是一个简单的数据加载和预处理示例：

import numpy as np
import pandas as pd
import torch

# 加载交通流数据
data = pd.read_csv('traffic_flow_data.csv')  # 假设数据存储在CSV文件中
data = data.values  # 转换为NumPy数组

# 划分训练集和测试集
train_ratio = 0.8
train_size = int(len(data) * train_ratio)
train_data = data[:train_size]
test_data = data[train_size:]

# 数据归一化
mean = np.mean(train_data)
std = np.std(train_data)
train_data = (train_data - mean) / std
test_data = (test_data - mean) / std

# 构建图邻接矩阵
# 这里假设你已经有了图的邻接矩阵，例如通过交通网络的拓扑结构得到
adj_matrix = np.load('adj_matrix.npy')  # 加载邻接矩阵
adj_matrix = torch.FloatTensor(adj_matrix)

# 生成训练和测试数据的输入输出序列
def generate_sequences(data, seq_length):
    inputs = []
    outputs = []
    for i in range(len(data) - seq_length):
        inputs.append(data[i:i+seq_length])
        outputs.append(data[i+seq_length])
    return np.array(inputs), np.array(outputs)

seq_length = 12  # 输入序列长度
train_inputs, train_outputs = generate_sequences(train_data, seq_length)
test_inputs, test_outputs = generate_sequences(test_data, seq_length)

train_inputs = torch.FloatTensor(train_inputs)
train_outputs = torch.FloatTensor(train_outputs)
test_inputs = torch.FloatTensor(test_inputs)
test_outputs = torch.FloatTensor(test_outputs)

2. 模型修改

确保现有的STGCN模型代码适用于交通流预测任务。以下是一个简单的STGCN模型示例：

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

class ChebConv(nn.Module):
    def __init__(self, in_channels, out_channels, K):
        super(ChebConv, self).__init__()
        self.K = K
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.weight = nn.Parameter(torch.FloatTensor(K, in_channels, out_channels))
        self.bias = nn.Parameter(torch.FloatTensor(out_channels))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weight)
        nn.init.zeros_(self.bias)

    def forward(self, x, L):
        batch_size, num_nodes, in_channels = x.size()
        outputs = []
        for b in range(batch_size):
            Tx_0 = x[b]
            output = torch.matmul(Tx_0, self.weight[0])
            if self.K > 1:
                Tx_1 = torch.matmul(L, Tx_0)
                output += torch.matmul(Tx_1, self.weight[1])
            for k in range(2, self.K):
                Tx_2 = 2 * torch.matmul(L, Tx_1) - Tx_0
                output += torch.matmul(Tx_2, self.weight[k])
                Tx_0, Tx_1 = Tx_1, Tx_2
            outputs.append(output)
        outputs = torch.stack(outputs, dim=0)
        outputs += self.bias
        return outputs

class STGCNBlock(nn.Module):
    def __init__(self, in_channels, spatial_channels, out_channels, num_nodes, K):
        super(STGCNBlock, self).__init__()
        self.temporal_conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 3), padding=(0, 1))
        self.cheb_conv = ChebConv(out_channels, spatial_channels, K)
        self.temporal_conv2 = nn.Conv2d(spatial_channels, out_channels, kernel_size=(1, 3), padding=(0, 1))
        self.batch_norm = nn.BatchNorm2d(num_nodes)

    def forward(self, x, L):
        x = self.temporal_conv1(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
        x = F.relu(x)
        batch_size, num_nodes, seq_length, in_channels = x.size()
        x = x.view(batch_size, num_nodes, -1)
        x = self.cheb_conv(x, L)
        x = F.relu(x)
        x = x.view(batch_size, num_nodes, seq_length, -1)
        x = self.temporal_conv2(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
        x = self.batch_norm(x)
        return x

class STGCN(nn.Module):
    def __init__(self, num_nodes, num_features, num_timesteps_input, num_timesteps_output, K):
        super(STGCN, self).__init__()
        self.block1 = STGCNBlock(in_channels=num_features, spatial_channels=16, out_channels=64, num_nodes=num_nodes, K=K)
        self.block2 = STGCNBlock(in_channels=64, spatial_channels=16, out_channels=64, num_nodes=num_nodes, K=K)
        self.fc = nn.Linear(num_timesteps_input * 64, num_timesteps_output)

    def forward(self, x, L):
        x = self.block1(x, L)
        x = self.block2(x, L)
        batch_size, num_nodes, seq_length, num_features = x.size()
        x = x.view(batch_size, num_nodes, -1)
        x = self.fc(x)
        return x

# 初始化模型
num_nodes = data.shape[1]
num_features = 1
num_timesteps_input = seq_length
num_timesteps_output = 1
K = 3  # Chebyshev多项式的阶数
model = STGCN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, K)

3. 训练和评估

训练模型并评估其在交通流预测上的性能。

import torch.optim as optim

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    outputs = model(train_inputs.unsqueeze(-1), adj_matrix)
    loss = criterion(outputs.squeeze(), train_outputs)
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 评估模型
model.eval()
with torch.no_grad():
    test_outputs_pred = model(test_inputs.unsqueeze(-1), adj_matrix)
    test_loss = criterion(test_outputs_pred.squeeze(), test_outputs)
    print(f'Test Loss: {test_loss.item():.4f}')

# 反归一化预测结果
test_outputs_pred = test_outputs_pred.squeeze().numpy() * std + mean
test_outputs = test_outputs.numpy() * std + mean