diff --git a/.vscode/launch.json b/.vscode/launch.json index 4b07128..12e148f 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -217,6 +217,14 @@ "program": "run.py", "console": "integratedTerminal", "args": "--config ./config/AEPSA/SolarEnergy.yaml" + }, + { + "name": "AEPSA_v2: METR-LA", + "type": "debugpy", + "request": "launch", + "program": "run.py", + "console": "integratedTerminal", + "args": "--config ./config/AEPSA/v2_METR-LA.yaml" } ] } \ No newline at end of file diff --git a/config/AEPSA/v2_METR-LA.yaml b/config/AEPSA/v2_METR-LA.yaml new file mode 100644 index 0000000..948ed6e --- /dev/null +++ b/config/AEPSA/v2_METR-LA.yaml @@ -0,0 +1,60 @@ +basic: + dataset: "METR-LA" + mode : "train" + device : "cuda:0" + model: "AEPSA_v2" + seed: 2023 + +data: + add_day_in_week: true + add_time_in_day: true + column_wise: false + days_per_week: 7 + default_graph: true + horizon: 24 + lag: 24 + normalizer: std + num_nodes: 207 + steps_per_day: 288 + test_ratio: 0.2 + tod: false + val_ratio: 0.2 + sample: 1 + input_dim: 1 + batch_size: 16 + +model: + pred_len: 24 + seq_len: 24 + patch_len: 6 + stride: 7 + dropout: 0.2 + gpt_layers: 9 + d_ff: 128 + gpt_path: ./GPT-2 + d_model: 64 + n_heads: 1 + input_dim: 1 + word_num: 1000 + num_nodes: 207 + +train: + batch_size: 16 + early_stop: true + early_stop_patience: 15 + epochs: 100 + grad_norm: false + loss_func: mae + lr_decay: true + lr_decay_rate: 0.3 + lr_decay_step: "5,20,40,70" + lr_init: 0.003 + max_grad_norm: 5 + real_value: true + weight_decay: 0 + debug: false + output_dim: 1 + log_step: 1000 + plot: false + mae_thresh: None + mape_thresh: 0.001 diff --git a/model/AEPSA/aepsav2.py b/model/AEPSA/aepsav2.py new file mode 100644 index 0000000..ceaa35d --- /dev/null +++ b/model/AEPSA/aepsav2.py @@ -0,0 +1,407 @@ +import torch +import torch.nn as nn +from transformers.models.gpt2.modeling_gpt2 import GPT2Model +from einops import rearrange +from model.AEPSA.normalizer import GumbelSoftmax +from model.AEPSA.reprogramming import ReprogrammingLayer +import torch.nn.functional as F + +# 该文件实现了基于动态图增强的时空序列预测模型 +# 主要包含三个类:DynamicGraphEnhancer(动态图增强器)、GraphEnhancedEncoder(图增强编码器)和AEPSA(主模型) +# 每个操作都标注了输入输出shape以帮助理解数据流向 + +class DynamicGraphEnhancer(nn.Module): + """ + 动态图增强器,基于节点嵌入自动生成图结构 + 参考DDGCRN的设计,使用节点嵌入和特征信息动态计算邻接矩阵 + """ + def __init__(self, num_nodes, in_dim, embed_dim=10): + # num_nodes: 节点数量 + # in_dim: 输入特征维度 + # embed_dim: 节点嵌入维度 + super().__init__() + self.num_nodes = num_nodes + self.embed_dim = embed_dim + + # 节点嵌入参数 [num_nodes, embed_dim] + self.node_embeddings = nn.Parameter( + torch.randn(num_nodes, embed_dim), requires_grad=True + ) + + # 特征转换层,用于生成动态调整的嵌入 + # 输入: [N, in_dim] -> 输出: [N, embed_dim] + self.feature_transform = nn.Sequential( + nn.Linear(in_dim, 16), # [N, in_dim] -> [N, 16] + nn.Sigmoid(), + nn.Linear(16, 2), # [N, 16] -> [N, 2] + nn.Sigmoid(), + nn.Linear(2, embed_dim) # [N, 2] -> [N, embed_dim] + ) + + # 注册单位矩阵作为固定的支持矩阵 [num_nodes, num_nodes] + self.register_buffer("eye", torch.eye(num_nodes)) + + def get_laplacian(self, graph, I, normalize=True): + """ + 计算归一化拉普拉斯矩阵 + + 参数: + graph: 邻接矩阵 [N, N] + I: 单位矩阵 [N, N] + normalize: 是否使用标准化拉普拉斯矩阵 + + 返回: + laplacian: 归一化拉普拉斯矩阵 [N, N] + """ + # 计算度矩阵的逆平方根 [N, N] + D_inv = torch.diag_embed(torch.sum(graph, -1) ** (-0.5)) # [N, N] + D_inv[torch.isinf(D_inv)] = 0.0 # 处理零除问题 + + if normalize: + # 归一化拉普拉斯矩阵: D^(-1/2) * graph * D^(-1/2) [N, N] + return torch.matmul(torch.matmul(D_inv, graph), D_inv) # [N, N] + else: + # 拉普拉斯矩阵: D^(-1/2) * (graph + I) * D^(-1/2) [N, N] + return torch.matmul(torch.matmul(D_inv, graph + I), D_inv) # [N, N] + + def forward(self, X): + """ + 参数: + X: 输入特征 [B, N, D],其中B为批次大小,N为节点数,D为特征维度 + + 返回: + laplacians: 动态生成的归一化拉普拉斯矩阵 [B, N, N] + """ + batch_size = X.size(0) + laplacians = [] + + # 获取单位矩阵 [N, N] + I = self.eye.to(X.device) + + for b in range(batch_size): + # 使用特征转换层生成动态嵌入调整因子 [N, embed_dim] + filt = self.feature_transform(X[b]) # [N, embed_dim] + + # 计算节点嵌入向量 [N, embed_dim] + nodevec = torch.tanh(self.node_embeddings * filt) # [N, embed_dim] + + # 通过节点嵌入的点积计算邻接矩阵 [N, N] + adj = F.relu(torch.matmul(nodevec, nodevec.transpose(0, 1))) # [N, N] + + # 计算归一化拉普拉斯矩阵 [N, N] + laplacian = self.get_laplacian(adj, I) # [N, N] + laplacians.append(laplacian) + + # 堆叠所有批次的拉普拉斯矩阵 [B, N, N] + return torch.stack(laplacians, dim=0) # [B, N, N] + +class GraphEnhancedEncoder(nn.Module): + """ + 基于Chebyshev多项式和动态拉普拉斯矩阵的图增强编码器 + 用于将动态图结构信息整合到特征编码中 + 优化版本:支持直接处理原始时间序列输入 + """ + def __init__(self, K=3, in_dim=64, hidden_dim=32, num_nodes=325, embed_dim=10, device='cpu', + temporal_dim=12, num_features=1): + # K: Chebyshev多项式阶数 + # in_dim: 输入特征维度 + # hidden_dim: 隐藏层维度 + # num_nodes: 节点数量 + # embed_dim: 嵌入维度 + # temporal_dim: 时间序列长度 + # num_features: 特征通道数量 + super().__init__() + self.K = K # Chebyshev多项式阶数 + self.in_dim = in_dim + self.hidden_dim = hidden_dim + self.device = device + self.temporal_dim = temporal_dim + self.num_features = num_features + + # 输入预处理层,适配正确的通道维度 + # 输入: [B, C, N, T] -> 输出: [B, in_dim, N, 1] + self.input_projection = nn.Sequential( + # 2D卷积:[B, C, N, T] -> [B, 16, N, T] + nn.Conv2d(num_features, 16, kernel_size=(1, 3), padding=(0, 1)), # [B, C, N, T] -> [B, 16, N, T] + nn.ReLU(), + # 2D卷积:[B, 16, N, T] -> [B, in_dim, N, 1],时间维度上的全局卷积 + nn.Conv2d(16, in_dim, kernel_size=(1, temporal_dim)), # [B, 16, N, T] -> [B, in_dim, N, 1] + nn.ReLU() + ) + + # 动态图增强器,用于生成动态拉普拉斯矩阵 + # 输入: [B, N, in_dim] -> 输出: [B, N, N] + self.graph_enhancer = DynamicGraphEnhancer(num_nodes, in_dim, embed_dim) + + # 谱系数 α_k (可学习参数) [K+1, 1] + self.alpha = nn.Parameter(torch.randn(K + 1, 1)) + + # 传播权重 W_k (可学习参数) + # 每个权重将Chebyshev多项式展开的结果从in_dim映射到hidden_dim + # 输入: [N, in_dim] -> 输出: [N, hidden_dim] + self.W = nn.ParameterList([ + nn.Parameter(torch.randn(in_dim, hidden_dim)) for _ in range(K + 1) + ]) + + self.to(device) + + def chebyshev_polynomials(self, L_tilde, X): + """ + 递归计算Chebyshev多项式展开 [T_0(L_tilde)X, ..., T_K(L_tilde)X] + + 参数: + L_tilde: 归一化拉普拉斯矩阵 [N, N] + X: 输入特征 [N, in_dim] + + 返回: + T_k_list: Chebyshev多项式展开列表 [K+1, N, in_dim] + """ + # T_0(X) = X [N, in_dim] + T_k_list = [X] + if self.K >= 1: + # T_1(X) = L_tilde * X [N, in_dim] + T_k_list.append(torch.matmul(L_tilde, X)) + for k in range(2, self.K + 1): + # T_k(X) = 2*L_tilde*T_{k-1}(X) - T_{k-2}(X) [N, in_dim] + T_k_list.append(2 * torch.matmul(L_tilde, T_k_list[-1]) - T_k_list[-2]) + # 返回Chebyshev多项式展开列表 [K+1, N, in_dim] + return T_k_list + + def forward(self, X): + """ + 参数: + X: 输入特征 [B, N, C, T] 或 [B, N, T](单特征情况) + B: 批次大小, N: 节点数, C: 特征通道数, T: 时间序列长度 + + 返回: + 增强后的特征 [B, N, hidden_dim*(K+1)] + """ + batch_size = X.size(0) + num_nodes = X.size(1) + + # 处理不同维度的输入 + if len(X.shape) == 4: # [B, N, C, T] + # 输入: [B, N, C, T] -> 输出: [B, C, N, T] + # 将输入转换为[B, C, N, T]格式,适合2D卷积 + x = X.permute(0, 2, 1, 3) # [B, C, N, T] + else: # [B, N, T] + # 输入: [B, N, T] -> 输出: [B, 1, N, T] + # 添加通道维度 + x = X.unsqueeze(1) # [B, 1, N, T] + + # 使用卷积投影层处理时间维度 + # 输入: [B, C, N, T] -> 输出: [B, in_dim, N, 1] + x_proj = self.input_projection(x) + # 输入: [B, in_dim, N, 1] -> 输出: [B, in_dim, N] + x_proj = x_proj.squeeze(-1) # [B, in_dim, N] + # 输入: [B, in_dim, N] -> 输出: [B, N, in_dim] + x_proj = x_proj.permute(0, 2, 1) # [B, N, in_dim] + + enhanced_features = [] + + # 动态生成拉普拉斯矩阵 + # 输入: [B, N, in_dim] -> 输出: [B, N, N] + laplacians = self.graph_enhancer(x_proj) # [B, N, N] + + for b in range(batch_size): + # 获取当前批次的拉普拉斯矩阵 [N, N] + L = laplacians[b] # [N, N] + + # 特征值缩放 + try: + # 计算最大特征值 [1] + lambda_max = torch.linalg.eigvalsh(L).max().real # [1] + # 避免除零问题 + if lambda_max < 1e-6: + lambda_max = 1.0 + # 缩放拉普拉斯矩阵到[-1, 1]区间 [N, N] + L_tilde = (2.0 / lambda_max) * L - torch.eye(L.size(0), device=L.device) # [N, N] + except: + # 如果计算特征值失败,使用单位矩阵 [N, N] + L_tilde = torch.eye(num_nodes, device=X.device) # [N, N] + + # 计算Chebyshev多项式展开 + # 输入: L_tilde [N, N], x_proj [N, in_dim] -> 输出: [K+1, N, in_dim] + T_k_list = self.chebyshev_polynomials(L_tilde, x_proj[b]) # [K+1, N, in_dim] + H_list = [] + + # 应用传播权重 + for k in range(self.K + 1): + # 矩阵乘法: [N, in_dim] × [in_dim, hidden_dim] -> [N, hidden_dim] + H_k = torch.matmul(T_k_list[k], self.W[k]) # [N, hidden_dim] + H_list.append(H_k) + + # 拼接所有尺度的特征 + # 输入: [K+1, N, hidden_dim] -> 输出: [N, hidden_dim*(K+1)] + X_enhanced = torch.cat(H_list, dim=-1) # [N, hidden_dim*(K+1)] + enhanced_features.append(X_enhanced) + + # 堆叠所有批次的增强特征 + # 输入: [B, N, hidden_dim*(K+1)] -> 输出: [B, N, hidden_dim*(K+1)] + return torch.stack(enhanced_features, dim=0) # [B, N, hidden_dim*(K+1)] + +class AEPSA(nn.Module): + """ + 自适应特征投影时空自注意力模型(AEPSA) + 整合动态图增强和预训练语言模型进行时空序列预测 + """ + def __init__(self, configs): + # configs: 包含模型所有配置的字典 + # 主要配置参数说明: + # device: 运行设备 + # pred_len: 预测序列长度 + # seq_len: 输入序列长度 + # patch_len: 补丁长度(已移除对应组件) + # input_dim: 输入特征维度 + # stride: 步长(已移除对应组件) + # dropout: Dropout概率 + # gpt_layers: 使用的GPT2层数 + # d_ff: 前馈网络隐藏层维度 + # gpt_path: 预训练GPT2模型路径 + # num_nodes: 节点数量 + # word_num: GumbelSoftmax词汇数量 + # d_model: 模型维度 + # n_heads: 注意力头数量 + # chebyshev_order: Chebyshev多项式阶数 + # graph_hidden_dim: 图编码器隐藏层维度 + # graph_embed_dim: 图编码器嵌入维度 + super(AEPSA, self).__init__() + self.device = configs['device'] + self.pred_len = configs['pred_len'] + self.seq_len = configs['seq_len'] + self.patch_len = configs['patch_len'] + self.input_dim = configs['input_dim'] + self.stride = configs['stride'] + self.dropout = configs['dropout'] + self.gpt_layers = configs['gpt_layers'] + self.d_ff = configs['d_ff'] + self.gpt_path = configs['gpt_path'] + self.num_nodes = configs.get('num_nodes', 325) # 添加节点数量配置 + + # GumbelSoftmax层,用于词汇选择 + # 输入: [vocab_size] -> 输出: [vocab_size](one-hot近似分布) + self.word_choice = GumbelSoftmax(configs['word_num']) + + self.d_model = configs['d_model'] + self.n_heads = configs['n_heads'] + self.d_keys = None + self.d_llm = 768 # GPT2隐藏层维度 + + self.patch_nums = int((self.seq_len - self.patch_len) / self.stride + 2) + self.head_nf = self.d_ff * self.patch_nums + + # 移除不再使用的patch_embedding层 + + # GPT2初始化 + # 加载预训练GPT2模型,输出注意力权重和隐藏状态 + self.gpts = GPT2Model.from_pretrained(self.gpt_path, output_attentions=True, output_hidden_states=True) + self.gpts.h = self.gpts.h[:self.gpt_layers] # 截取指定层数 + self.gpts.apply(self.reset_parameters) + + # 获取GPT2词嵌入权重 + # 形状: [vocab_size, d_llm] + self.word_embeddings = self.gpts.get_input_embeddings().weight.to(self.device) + self.vocab_size = self.word_embeddings.shape[0] + # 映射层,将词汇表维度映射到1维 + # 输入: [vocab_size] -> 输出: [1] + self.mapping_layer = nn.Linear(self.vocab_size, 1) + # 重编程层,用于特征映射和注意力计算 + # 输入: [B, N, d_model], [d_llm], [d_llm] -> 输出: [B, N, d_model] + self.reprogramming_layer = ReprogrammingLayer(self.d_model, self.n_heads, self.d_keys, self.d_llm) + + # 动态图增强编码器 + # 输入: [B, N, C, T] -> 输出: [B, N, hidden_dim*(K+1)] + self.graph_encoder = GraphEnhancedEncoder( + K=configs.get('chebyshev_order', 3), # Chebyshev多项式阶数 + in_dim=self.d_model, # 输入特征维度 + hidden_dim=configs.get('graph_hidden_dim', 32), # 隐藏层维度 + num_nodes=self.num_nodes, # 节点数量 + embed_dim=configs.get('graph_embed_dim', 10), # 节点嵌入维度 + device=self.device, # 运行设备 + temporal_dim=self.seq_len, # 时间序列长度 + num_features=self.input_dim # 特征通道数 + ) + + # 图特征投影层,将图增强特征维度转换为d_model + # 输入: [B, N, hidden_dim*(K+1)] -> 输出: [B, N, d_model] + self.graph_projection = nn.Linear( + configs.get('graph_hidden_dim', 32) * (configs.get('chebyshev_order', 3) + 1), + self.d_model + ) + + self.out_mlp = nn.Sequential( + nn.Linear(self.d_llm, 128), + nn.ReLU(), + nn.Linear(128, self.pred_len) + ) + + for i, (name, param) in enumerate(self.gpts.named_parameters()): + if 'wpe' in name: + param.requires_grad = True + else: + param.requires_grad = False + + def reset_parameters(self, module): + if hasattr(module, 'weight') and module.weight is not None: + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if hasattr(module, 'bias') and module.bias is not None: + torch.nn.init.zeros_(module.bias) + + def forward(self, x): + """ + 前向传播函数 + 输入: + x: 输入数据 [B, T, N, C],其中B为批次大小,T为时间步长,N为节点数,C为特征通道数 + + 返回: + outputs: 预测结果 [B, pred_len, N, 1] + """ + # 只保留第一个特征通道 + # 输入: [B, T, N, C] -> 输出: [B, T, N, 1] + x = x[..., :1] # [B, T, N, 1] + + # 调整输入维度以适配图编码器 + # 输入: [B, T, N, 1] -> 输出: [B, N, 1, T] + x_enc = rearrange(x, 'b t n c -> b n c t') # [B, N, 1, T] + + # 应用图增强编码器获取增强特征 + # 输入: [B, N, 1, T] -> 输出: [B, N, hidden_dim*(K+1)] + graph_enhanced = self.graph_encoder(x_enc) # [B, N, hidden_dim*(K+1)] + + # 投影图增强特征到模型维度 + # 输入: [B, N, hidden_dim*(K+1)] -> 输出: [B, N, d_model] + enc_out = self.graph_projection(graph_enhanced) # [B, N, d_model] + + # 处理词嵌入权重,为注意力机制准备 + # 输入: [vocab_size, d_llm] -> 输出: [d_llm, vocab_size] -> [d_llm, vocab_size] + self.mapping_layer(self.word_embeddings.permute(1, 0)).permute(1, 0) # [vocab_size, d_llm] + + # 使用GumbelSoftmax选择词汇 + # 输入: [d_llm, 1] -> 输出: [d_llm, 1] + masks = self.word_choice(self.mapping_layer.weight.data.permute(1,0)) # [d_llm, 1] + + # 获取选中的源嵌入 + # 输入: [vocab_size, d_llm] 与 masks -> 输出: [selected_words, d_llm] + source_embeddings = self.word_embeddings[masks==1] # [selected_words, d_llm] + + # 应用重编程层处理特征和源嵌入 + # 输入: [B, N, d_model], [selected_words, d_llm], [selected_words, d_llm] -> 输出: [B, N, d_model] + enc_out = self.reprogramming_layer(enc_out, source_embeddings, source_embeddings) # [B, N, d_model] + + # 通过GPT2模型处理增强特征 + # 输入: [B, N, d_model] -> 输出: [B, N, d_llm] + enc_out = self.gpts(inputs_embeds=enc_out).last_hidden_state # [B, N, d_llm] + + # 使用MLP预测未来时间步 + # 输入: [B, N, d_llm] -> 输出: [B, N, pred_len] + dec_out = self.out_mlp(enc_out) # [B, N, pred_len] + + # 添加通道维度 + # 输入: [B, N, pred_len] -> 输出: [B, N, pred_len, 1] + outputs = dec_out.unsqueeze(dim=-1) # [B, N, pred_len, 1] + + # 调整维度顺序为 [B, pred_len, N, 1] + # 输入: [B, N, pred_len, 1] -> 输出: [B, pred_len, N, 1] + outputs = outputs.permute(0, 2, 1, 3) # [B, pred_len, N, 1] + + return outputs # [B, pred_len, N, 1] \ No newline at end of file diff --git a/model/model_selector.py b/model/model_selector.py index 5043202..c669d82 100755 --- a/model/model_selector.py +++ b/model/model_selector.py @@ -24,6 +24,8 @@ from model.STGNRDE.Make_model import make_model as make_nrde_model from model.STAWnet.STAWnet import STAWnet from model.REPST.repst import repst as REPST from model.AEPSA.aepsa import AEPSA as AEPSA +from model.AEPSA.aepsav2 import AEPSA as AEPSAv2 + def model_selector(config): @@ -82,3 +84,5 @@ def model_selector(config): return REPST(model_config) case "AEPSA": return AEPSA(model_config) + case "AEPSA_v2": + return AEPSAv2(model_config)