python import numpy as np import torch from torch.utils.data import DataLoader, TensorDataset

def create_sequences(data, seq_len, pred_len): x, y = [], [] for i in range(len(data) - seq_len - pred_len + 1): x.append(data[:, i:i+seq_len]) y.append(data[:, i+seq_len:i+seq_len+pred_len]) return torch.FloatTensor(np.array(x)), torch.FloatTensor(np.array(y))

seq_len = 12 # 用过去12个时刻预测 pred_len = 3 # 预测未来3个时刻

假设 traffic: (num_nodes, total_timesteps)

x, y = create_sequences(traffic, seq_len, pred_len) dataset = TensorDataset(x.permute(0, 2, 1), y.permute(0, 2, 1)) # (batch, seq, nodes)

### 4.2 构建图邻接矩阵

```python
import scipy.sparse as sp

# 根据距离或路网构造加权邻接矩阵 w
def normalize_adj(adj):
    adj = sp.coo_matrix(adj)
    rowsum = np.array(adj.sum(1))
    d_inv_sqrt = np.power(rowsum, -0.5).flatten()
    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
    return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()

adj = normalize_adj(weighted_adjacency)
adj = torch.FloatTensor(adj.todense())  # (N, N)

4.3 定义模型：空间卷积 + 时间卷积

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

class ChebConv(nn.Module):
    # K阶切比雪夫图卷积
    def __init__(self, in_channels, out_channels, K, cheb_polynomials):
        super().__init__()
        self.K = K
        self.cheb_polynomials = cheb_polynomials
        self.weight = nn.Parameter(torch.Tensor(K, in_channels, out_channels))
        self.bias = nn.Parameter(torch.Tensor(out_channels))
        
    def forward(self, x):
        # x: (batch, nodes, in_channels)
        batch, N, _ = x.shape
        output = torch.zeros(batch, N, self.weight.shape[-1]).to(x.device)
        for k in range(self.K):
            T_k = self.cheb_polynomials[k].to(x.device)  # (N, N)
            rhs = torch.bmm(T_k.unsqueeze(0).expand(batch, -1, -1), x)
            output += rhs @ self.weight[k]
        return output + self.bias

class TemporalConvLayer(nn.Module):
    # 门控一维卷积
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, 2*out_channels, kernel_size, padding='same')
        
    def forward(self, x):
        # x: (batch, nodes, time, channels)
        b, n, t, c = x.shape
        x = x.permute(0, 1, 3, 2).reshape(b*n, c, t)
        out = self.conv(x)  # (b*n, 2*out_c, t)
        P, Q = out.chunk(2, dim=1)
        return (P + Q.sigmoid() * P).reshape(b, n, -1, t).permute(0, 1, 3, 2)