上传文件至 'models/pressure-predictor/90天TMP预测模型源码修改'

models/pressure-predictor/90天TMP预测模型源码修改/gat_lstm.py

@@ -0,0 +1,102 @@
+# gat_lstm.py
+import torch
+import torch.nn as nn    # PyTorch神经网络模块
+
+# 单个独立模型（对应1个因变量）
+class SingleGATLSTM(nn.Module):
+    def __init__(self, args):
+        """
+        单个子模型：包含GAT-LSTM层和输出层，用于预测1个目标指标
+        参数：
+            args: 配置参数（含特征数、隐藏层大小等）
+        """
+        super(SingleGATLSTM, self).__init__()
+        self.args = args
+        
+        # 独立的LSTM层
+        self.lstm = nn.LSTM(
+            input_size=args.feature_num,
+            hidden_size=args.hidden_size,
+            num_layers=args.num_layers,
+            batch_first=True
+        )
+        
+        # 独立的输出层
+        self.final_linear = nn.Sequential(
+            nn.Linear(args.hidden_size, args.hidden_size),
+            nn.LeakyReLU(0.01),
+            nn.Dropout(args.dropout * 0.4),
+            nn.Linear(args.hidden_size, args.output_size)
+        )
+        
+        self._init_weights()
+        
+    def _init_weights(self):
+        """初始化网络权重，加速模型收敛"""
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # 初始化LSTM权重
+        for name, param in self.lstm.named_parameters():
+            if 'weight_ih' in name:
+                nn.init.xavier_uniform_(param.data)
+            elif 'weight_hh' in name:
+                nn.init.orthogonal_(param.data)
+            elif 'bias' in name:
+                param.data.fill_(0)
+                n = param.size(0)
+                start, end = n // 4, n // 2
+                param.data[start:end].fill_(1)
+        
+    def forward(self, x):
+        """
+        前向传播：输入序列经过LSTM和输出层，得到预测结果
+        参数：
+            x: 输入序列，形状为[batch_size, seq_len, feature_num]
+        返回：
+            output: 预测结果，形状为[batch_size, output_size]
+        """
+        batch_size, seq_len, feature_num = x.size()
+        lstm_out, _ = self.lstm(x)
+        # 取最后一个时间步的输出
+        last_out = lstm_out[:, -1, :]
+        
+        # 输出层预测
+        output = self.final_linear(last_out)
+        return output  # [batch_size, output_size]
+
+
+# 16个独立模型的容器（总模型）
+class GAT_LSTM(nn.Module):
+    def __init__(self, args):
+        """
+        总模型：包含多个SingleGATLSTM子模型，分别预测不同的目标
+        参数：
+            args: 配置参数（含labels_num，即子模型数量）
+        """
+        super(GAT_LSTM, self).__init__()
+        self.args = args
+        # 创建16个独立模型（数量由labels_num指定）
+        self.models = nn.ModuleList([SingleGATLSTM(args) for _ in range(args.labels_num)])
+    
+    def set_edge_index(self, edge_index):
+        self.edge_index = edge_index  # 将传入的edge_index保存到模型内部
+        
+    def forward(self, x):
+        """
+        前向传播：所有子模型并行处理输入，拼接预测结果
+        参数：
+            x: 输入序列，形状为[batch_size, seq_len, feature_num]
+        返回：
+            拼接后的预测结果，形状为[batch_size, output_size * labels_num]
+        """
+        outputs = []
+        for model in self.models:
+            outputs.append(model(x))  # 每个输出为[batch, output_size]
+        return torch.cat(outputs, dim=1)  # 拼接后[batch, output_size * labels_num]

models/pressure-predictor/90天TMP预测模型源码修改/main.py

@@ -0,0 +1,70 @@
+# main.py
+import os
+import torch
+import numpy as np
+import random
+from gat_lstm import GAT_LSTM
+from data_trainer import Trainer
+from args import lstm_args_parser
+from torch.nn import MSELoss
+from data_preprocessor import DataPreprocessor
+
+def set_seed(seed):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def main():
+    args = lstm_args_parser()
+    set_seed(args.random_seed)
+    
+    device = torch.device(f"cuda:{args.device}" if torch.cuda.is_available() else "cpu")
+    args.device = device  # 将device存入args，方便后续使用
+
+    # 数据预处理
+    data = DataPreprocessor.read_and_combine_csv_files(args)
+    train_loader, val_loader, test_loader, _ = DataPreprocessor.load_and_process_data(args, data)
+    
+    # 初始化包含16个子模型的整体模型
+    model = GAT_LSTM(args).to(device)
+
+    # 初始化训练器和MSE损失函数
+    trainer = Trainer(model, args, data)
+    criterion = MSELoss()
+
+    # 优化器：优化所有子模型的参数（联合训练）
+    optimizer = torch.optim.Adam(
+        model.parameters(),  # 整体模型参数
+        lr=args.lr
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=args.scheduler_step_size,
+        gamma=args.scheduler_gamma
+    )
+
+    # 整体训练大模型（包含所有8/16个子模型）
+    print("=== 开始训练包含8/16个子模型的整体模型 ===")
+    trainer.train_full_model(
+        train_loader,
+        val_loader,
+        optimizer,
+        criterion,
+        scheduler
+    )
+
+    # 保存包含所有8/16个子模型参数的整体模型
+    trainer.save_model()
+    print("\n=== 模型训练完成，已保存整体模型 ===")
+
+    # 评估模型
+    trainer.evaluate_model(test_loader, MSELoss())
+
+if __name__ == "__main__":
+    main()