优化trainer和run

.vscode/settings.json

@@ -1,5 +1,5 @@
 {
-    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultEnvManager": "ms-python.python:system",
-    "python-envs.defaultPackageManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:pip",
     "python-envs.pythonProjects": []
 }

config/REPST/PEMS-BAY.yaml

@@ -3,6 +3,7 @@ basic:
   mode : "train"
   device : "cuda:1"
   model: "REPST"
+  seed: 2023
 
 data:
   add_day_in_week: true
@@ -49,7 +50,6 @@ train:
   lr_init: 0.003
   max_grad_norm: 5
   real_value: true
-  seed: 12
   weight_decay: 0
   debug: false
   output_dim: 1

dataloader/PeMSDdataloader.py

@@ -2,112 +2,88 @@ from utils.normalization import normalize_dataset
 from dataloader.data_selector import load_st_dataset
 
 import numpy as np
-import gc
 import torch
 
 
-
 def get_dataloader(args, normalizer="std", single=True):
-    data = load_st_dataset(args)  # 加载数据
+    data = load_st_dataset(args)
     args = args["data"]
-    L, N, F = data.shape  # 数据形状
+    L, N, F = data.shape
 
-    # Step 1: data -> x,y
+    # Generate sliding windows for main data and add time features
+    x, y = _prepare_data_with_windows(data, args, single)
+    
+    # Split data
+    split_fn = split_data_by_days if args["test_ratio"] > 1 else split_data_by_ratio
+    x_train, x_val, x_test = split_fn(x, args["val_ratio"], args["test_ratio"])
+    y_train, y_val, y_test = split_fn(y, args["val_ratio"], args["test_ratio"])
+
+    # Normalize x and y using the same scaler
+    scaler = _normalize_data(x_train, x_val, x_test, args, normalizer)
+    _apply_existing_scaler(y_train, y_val, y_test, scaler, args)
+
+    # Create dataloaders
+    return (
+        _create_dataloader(x_train, y_train, args["batch_size"], True, False),
+        _create_dataloader(x_val, y_val, args["batch_size"], False, False),
+        _create_dataloader(x_test, y_test, args["batch_size"], False, False),
+        scaler
+    )
+
+
+def _prepare_data_with_windows(data, args, single):
+    # Generate sliding windows for main data
     x = add_window_x(data, args["lag"], args["horizon"], single)
     y = add_window_y(data, args["lag"], args["horizon"], single)
 
-    del data
+    # Generate time features
-    gc.collect()
+    time_features = _generate_time_features(data.shape[0], args)
     
-    # Step 2: time_in_day, day_in_week -> day, week
+    # Add time features to x and y
+    x = _add_time_features(x, time_features, args["lag"], args["horizon"], single, add_window_x)
+    y = _add_time_features(y, time_features, args["lag"], args["horizon"], single, add_window_y)
+    
+    return x, y
+
+
+def _generate_time_features(L, args):
+    N = args["num_nodes"]
     time_in_day = [i % args["steps_per_day"] / args["steps_per_day"] for i in range(L)]
     time_in_day = np.tile(np.array(time_in_day), [1, N, 1]).transpose((2, 1, 0))
-    day_in_week = [
+    
-        (i // args["steps_per_day"]) % args["days_per_week"] for i in range(L)
+    day_in_week = [(i // args["steps_per_day"]) % args["days_per_week"] for i in range(L)]
-    ]
     day_in_week = np.tile(np.array(day_in_week), [1, N, 1]).transpose((2, 1, 0))
     
-    x_day = add_window_x(time_in_day, args["lag"], args["horizon"], single)
+    return time_in_day, day_in_week
-    x_week = add_window_x(day_in_week, args["lag"], args["horizon"], single)
 
-    # Step 3 day, week, x, y --> x, y
-    x = np.concatenate([x, x_day, x_week], axis=-1)
 
-    del x_day, x_week
+def _add_time_features(data, time_features, lag, horizon, single, window_fn):
-    gc.collect()
+    time_in_day, day_in_week = time_features
+    time_day = window_fn(time_in_day, lag, horizon, single)
+    time_week = window_fn(day_in_week, lag, horizon, single)
+    return np.concatenate([data, time_day, time_week], axis=-1)
 
-    # Step 4 x,y --> x_train, x_val, x_test, y_train, y_val, y_test
-    if args["test_ratio"] > 1:
-        x_train, x_val, x_test = split_data_by_days(
-            x, args["val_ratio"], args["test_ratio"]
-        )
-    else:
-        x_train, x_val, x_test = split_data_by_ratio(
-            x, args["val_ratio"], args["test_ratio"]
-        )
 
-    del x
+def _normalize_data(train_data, val_data, test_data, args, normalizer):
-    gc.collect()
+    scaler = normalize_dataset(train_data[..., : args["input_dim"]], normalizer, args["column_wise"])
     
-    # Normalization
+    for data in [train_data, val_data, test_data]:
-    scaler = normalize_dataset(
+        data[..., : args["input_dim"]] = scaler.transform(data[..., : args["input_dim"]])
-        x_train[..., : args["input_dim"]], normalizer, args["column_wise"]
-    )
-    x_train[..., : args["input_dim"]] = scaler.transform(
-        x_train[..., : args["input_dim"]]
-    )
-    x_val[..., : args["input_dim"]] = scaler.transform(x_val[..., : args["input_dim"]])
-    x_test[..., : args["input_dim"]] = scaler.transform(
-        x_test[..., : args["input_dim"]]
-    )
     
-    y_day = add_window_y(time_in_day, args["lag"], args["horizon"], single)
+    return scaler
-    y_week = add_window_y(day_in_week, args["lag"], args["horizon"], single)
 
-    del time_in_day, day_in_week
-    gc.collect()
 
-    y = np.concatenate([y, y_day, y_week], axis=-1)
+def _apply_existing_scaler(train_data, val_data, test_data, scaler, args):
+    for data in [train_data, val_data, test_data]:
+        data[..., : args["input_dim"]] = scaler.transform(data[..., : args["input_dim"]])
 
-    del y_day, y_week
-    gc.collect()
 
-    # Split Y
+def _create_dataloader(X_data, Y_data, batch_size, shuffle, drop_last):
-    if args["test_ratio"] > 1:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        y_train, y_val, y_test = split_data_by_days(
+    X_tensor = torch.tensor(X_data, dtype=torch.float32, device=device)
-            y, args["val_ratio"], args["test_ratio"]
+    Y_tensor = torch.tensor(Y_data, dtype=torch.float32, device=device)
-        )
+    dataset = torch.utils.data.TensorDataset(X_tensor, Y_tensor)
-    else:
+    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
-        y_train, y_val, y_test = split_data_by_ratio(
-            y, args["val_ratio"], args["test_ratio"]
-        )
-
-    del y
-    gc.collect()
-
-    # Step 5: x_train y_train x_val y_val x_test y_test --> train val test
-    train_dataloader = data_loader(
-        x_train, y_train, args["batch_size"], shuffle=True, drop_last=True
-    )
-
-    del x_train, y_train
-    gc.collect()
-
-    val_dataloader = data_loader(
-        x_val, y_val, args["batch_size"], shuffle=False, drop_last=True
-    )
-
-    del x_val, y_val
-    gc.collect()
-
-    test_dataloader = data_loader(
-        x_test, y_test, args["batch_size"], shuffle=False, drop_last=False
-    )
-
-    del x_test, y_test
-    gc.collect()
-
-    return train_dataloader, val_dataloader, test_dataloader, scaler
 
 
 def split_data_by_days(data, val_days, test_days, interval=30):
@@ -128,17 +104,29 @@ def split_data_by_ratio(data, val_ratio, test_ratio):
     return train_data, val_data, test_data
 
 
-def data_loader(X, Y, batch_size, shuffle=True, drop_last=True):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    X = torch.tensor(X, dtype=torch.float32, device=device)
-    Y = torch.tensor(Y, dtype=torch.float32, device=device)
-    data = torch.utils.data.TensorDataset(X, Y)
-    dataloader = torch.utils.data.DataLoader(
-        data, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last
-    )
-    return dataloader
 
 
+def _generate_windows(data, window=3, horizon=1, offset=0):
+    """
+    Internal helper function to generate sliding windows.
+    
+    :param data: Input data
+    :param window: Window size
+    :param horizon: Horizon size
+    :param offset: Offset from window start
+    :return: Windowed data
+    """
+    length = len(data)
+    end_index = length - horizon - window + 1
+    windows = []
+    index = 0
+
+    while index < end_index:
+        windows.append(data[index + offset : index + offset + window])
+        index += 1
+
+    return np.array(windows)
+
 def add_window_x(data, window=3, horizon=1, single=False):
     """
     Generate windowed X values from the input data.
@@ -149,17 +137,7 @@ def add_window_x(data, window=3, horizon=1, single=False):
     :param single: If True, generate single-step windows, else multi-step
     :return: X with shape [B, W, ...]
     """
-    length = len(data)
+    return _generate_windows(data, window, horizon, offset=0)
-    end_index = length - horizon - window + 1
-    x = []  # Sliding windows
-    index = 0
-
-    while index < end_index:
-        x.append(data[index : index + window])
-        index += 1
-
-    return np.array(x)
-
 
 def add_window_y(data, window=3, horizon=1, single=False):
     """
@@ -171,21 +149,10 @@ def add_window_y(data, window=3, horizon=1, single=False):
     :param single: If True, generate single-step windows, else multi-step
     :return: Y with shape [B, H, ...]
     """
-    length = len(data)
+    offset = window if not single else window + horizon - 1
-    end_index = length - horizon - window + 1
+    return _generate_windows(data, window=1 if single else horizon, horizon=horizon, offset=offset)
-    y = []  # Horizon values
-    index = 0
-
-    while index < end_index:
-        if single:
-            y.append(data[index + window + horizon - 1 : index + window + horizon])
-        else:
-            y.append(data[index + window : index + window + horizon])
-        index += 1
-
-    return np.array(y)
-
 
 if __name__ == "__main__":
-    res = load_st_dataset("SD", 1)
+    from dataloader.data_selector import load_st_dataset
-    k = 1
+    res = load_st_dataset({"dataset": "SD"})
+    print(f"Dataset shape: {res.shape}")

model/AEPSA/repst.py

@@ -1,4 +1,3 @@
-from tkinter import Y
 import torch
 import torch.nn as nn
 from transformers.models.gpt2.modeling_gpt2 import GPT2Model

model/REPST/reprogramming.py

@@ -30,7 +30,6 @@ class TokenEmbedding(nn.Module):
 
     def forward(self, x):
         b, n, m, pn, pl = x.shape  # batch, node, feature, patch_num, patch_len
-        # 为什么没permute后reshape?
         x = x.permute(0, 1, 4, 3, 2)
         x = self.tokenConv(x.reshape(b*n, pl, m*pn))  # batch*node, patch_len, feature*patch_num
         x = self.confusion_layer(x)

model/REPST/repst.py

@@ -1,4 +1,3 @@
-from tkinter import Y
 import torch
 import torch.nn as nn
 from transformers.models.gpt2.modeling_gpt2 import GPT2Model

run.py

@@ -13,7 +13,7 @@ from trainer.trainer_selector import select_trainer
 def main():
     args = parse_args()
     args = init.init_device(args)
-    init.init_seed(args["train"]["seed"])
+    init.init_seed(args["basic"]["seed"])
     model = init.init_model(args)
 
     # Load dataset

trainer/Trainer.py

@@ -10,11 +10,14 @@ from tqdm import tqdm
 
 
 class TrainingStats:
+    """记录训练过程中的统计信息"""
+    
     def __init__(self, device):
         self.device = device
         self.reset()
 
     def reset(self):
+        """重置所有统计数据"""
         self.gpu_mem_usage_list = []
         self.cpu_mem_usage_list = []
         self.train_time_list = []
@@ -24,9 +27,11 @@ class TrainingStats:
         self.end_time = None
 
     def start_training(self):
+        """记录训练开始时间"""
         self.start_time = time.time()
 
     def end_training(self):
+        """记录训练结束时间"""
         self.end_time = time.time()
 
     def record_step_time(self, duration, mode):
@@ -51,6 +56,10 @@ class TrainingStats:
         self.cpu_mem_usage_list.append(cpu_mem)
         self.gpu_mem_usage_list.append(gpu_mem)
 
+    def _calculate_average(self, values_list):
+        """安全计算平均值，避免除零错误"""
+        return sum(values_list) / len(values_list) if values_list else 0
+
     def report(self, logger):
         """在训练结束时输出汇总统计"""
         if not self.start_time or not self.end_time:
@@ -58,26 +67,10 @@ class TrainingStats:
             return
 
         total_time = self.end_time - self.start_time
-        avg_gpu_mem = (
+        avg_gpu_mem = self._calculate_average(self.gpu_mem_usage_list)
-            sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list)
+        avg_cpu_mem = self._calculate_average(self.cpu_mem_usage_list)
-            if self.gpu_mem_usage_list
+        avg_train_time = self._calculate_average(self.train_time_list)
-            else 0
+        avg_infer_time = self._calculate_average(self.infer_time_list)
-        )
-        avg_cpu_mem = (
-            sum(self.cpu_mem_usage_list) / len(self.cpu_mem_usage_list)
-            if self.cpu_mem_usage_list
-            else 0
-        )
-        avg_train_time = (
-            sum(self.train_time_list) / len(self.train_time_list)
-            if self.train_time_list
-            else 0
-        )
-        avg_infer_time = (
-            sum(self.infer_time_list) / len(self.infer_time_list)
-            if self.infer_time_list
-            else 0
-        )
         iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
 
         logger.info("===== Training Summary =====")
@@ -93,6 +86,8 @@ class TrainingStats:
 
 
 class Trainer:
+    """模型训练器，负责整个训练流程的管理"""
+    
     def __init__(
         self,
         model,
@@ -105,26 +100,42 @@ class Trainer:
         args,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
         self.device = args["basic"]["device"]
-        args = args["train"]
+        train_args = args["train"]
+        
+        # 模型和训练相关组件
         self.model = model
         self.loss = loss
         self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        
+        # 数据加载器
         self.train_loader = train_loader
         self.val_loader = val_loader
         self.test_loader = test_loader
+        
+        # 数据处理工具
         self.scaler = scaler
-        self.args = args
+        self.args = train_args
-        self.lr_scheduler = lr_scheduler
+        
+        # 统计信息
         self.train_per_epoch = len(train_loader)
         self.val_per_epoch = len(val_loader) if val_loader else 0
 
-        # Paths for saving models and logs
+        # 初始化路径、日志和统计
+        self._initialize_paths(train_args)
+        self._initialize_logger(train_args)
+        self._initialize_stats()
+    
+    def _initialize_paths(self, args):
+        """初始化模型保存路径"""
         self.best_path = os.path.join(args["log_dir"], "best_model.pth")
         self.best_test_path = os.path.join(args["log_dir"], "best_test_model.pth")
         self.loss_figure_path = os.path.join(args["log_dir"], "loss.png")
     
-        # Initialize logger
+    def _initialize_logger(self, args):
+        """初始化日志记录器"""
         if not os.path.isdir(args["log_dir"]) and not args["debug"]:
             os.makedirs(args["log_dir"], exist_ok=True)
         self.logger = get_logger(
@@ -132,10 +143,13 @@ class Trainer:
         )
         self.logger.info(f"Experiment log path in: {args['log_dir']}")
     
-        # Stats tracker
+    def _initialize_stats(self):
+        """初始化统计信息记录器"""
         self.stats = TrainingStats(device=self.device)
 
     def _run_epoch(self, epoch, dataloader, mode):
+        """运行一个训练/验证/测试epoch"""
+        # 设置模型模式和是否进行优化
         if mode == "train":
             self.model.train()
             optimizer_step = True
@@ -143,48 +157,64 @@ class Trainer:
             self.model.eval()
             optimizer_step = False
 
+        # 初始化变量
         total_loss = 0
         epoch_time = time.time()
-
         y_pred, y_true = [], []
 
+        # 训练/验证循环
         with torch.set_grad_enabled(optimizer_step):
-            progress_bar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}")
+            progress_bar = tqdm(
-            for batch_idx, (data, target) in progress_bar:
+                enumerate(dataloader), 
+                total=len(dataloader), 
+                desc=f"{mode.capitalize()} Epoch {epoch}"
+            )
+            
+            for _, (data, target) in progress_bar:
+                # 记录步骤开始时间
                 start_time = time.time()
 
+                # 前向传播
                 label = target[..., : self.args["output_dim"]]
                 output = self.model(data).to(self.device)
-
-                if self.args["real_value"]:
-                    output = self.scaler.inverse_transform(output)
-
                 loss = self.loss(output, label)
+
+                # 反归一化
+                self.scaler.inverse_transform(output)
+                self.scaler.inverse_transform(label)
+
+                # 反向传播和优化（仅在训练模式）
                 if optimizer_step and self.optimizer is not None:
                     self.optimizer.zero_grad()
                     loss.backward()
 
+                    # 梯度裁剪（如果需要）
                     if self.args["grad_norm"]:
                         torch.nn.utils.clip_grad_norm_(
                             self.model.parameters(), self.args["max_grad_norm"]
                         )
                     self.optimizer.step()
 
+                # 记录步骤时间和内存使用
                 step_time = time.time() - start_time
                 self.stats.record_step_time(step_time, mode)
 
+                # 累积损失和预测结果
                 total_loss += loss.item()
                 y_pred.append(output.detach().cpu())
                 y_true.append(label.detach().cpu())
 
-                # Update progress bar with the current loss
+                # 更新进度条
                 progress_bar.set_postfix(loss=loss.item())
 
+        # 合并所有批次的预测结果
         y_pred = torch.cat(y_pred, dim=0)
         y_true = torch.cat(y_true, dim=0)
 
+        # 计算平均损失
         avg_loss = total_loss / len(dataloader)
-        # 输出指标
+        
+        # 计算并记录指标
         mae, rmse, mape = all_metrics(
             y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
         )
@@ -192,7 +222,7 @@ class Trainer:
             f"Epoch #{epoch:02d}: {mode.capitalize():<5} MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s"
         )
 
-        # 记录内存
+        # 记录内存使用情况
         self.stats.record_memory_usage()
 
         return avg_loss
@@ -207,22 +237,29 @@ class Trainer:
         return self._run_epoch(epoch, self.test_loader, "test")
 
     def train(self):
+        """执行完整的训练流程"""
+        # 初始化最佳模型和损失记录
         best_model, best_test_model = None, None
         best_loss, best_test_loss = float("inf"), float("inf")
         not_improved_count = 0
 
+        # 开始训练
         self.stats.start_training()
         self.logger.info("Training process started")
 
+        # 训练循环
         for epoch in range(1, self.args["epochs"] + 1):
+            # 训练、验证和测试一个epoch
             train_epoch_loss = self.train_epoch(epoch)
             val_epoch_loss = self.val_epoch(epoch)
             test_epoch_loss = self.test_epoch(epoch)
 
+            # 检查梯度爆炸
             if train_epoch_loss > 1e6:
                 self.logger.warning("Gradient explosion detected. Ending...")
                 break
 
+            # 更新最佳验证模型
             if val_epoch_loss < best_loss:
                 best_loss = val_epoch_loss
                 not_improved_count = 0
@@ -231,6 +268,31 @@ class Trainer:
             else:
                 not_improved_count += 1
 
+            # 检查早停条件
+            if self._should_early_stop(not_improved_count):
+                break
+
+            # 更新最佳测试模型
+            if test_epoch_loss < best_test_loss:
+                best_test_loss = test_epoch_loss
+                best_test_model = copy.deepcopy(self.model.state_dict())
+
+        # 保存最佳模型
+        if not self.args["debug"]:
+            self._save_best_models(best_model, best_test_model)
+
+        # 结束训练并输出统计信息
+        self.stats.end_training()
+        self.stats.report(self.logger)
+
+        # 最终评估
+        self._finalize_training(best_model, best_test_model)
+
+        # 输出模型参数量
+        self._log_model_params()
+    
+    def _should_early_stop(self, not_improved_count):
+        """检查是否满足早停条件"""
         if (
             self.args["early_stop"]
             and not_improved_count == self.args["early_stop_patience"]
@@ -238,25 +300,19 @@ class Trainer:
             self.logger.info(
                 f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops."
             )
-                break
+            return True
+        return False
     
-            if test_epoch_loss < best_test_loss:
+    def _save_best_models(self, best_model, best_test_model):
-                best_test_loss = test_epoch_loss
+        """保存最佳模型到文件"""
-                best_test_model = copy.deepcopy(self.model.state_dict())
-
-        if not self.args["debug"]:
         torch.save(best_model, self.best_path)
         torch.save(best_test_model, self.best_test_path)
         self.logger.info(
             f"Best models saved at {self.best_path} and {self.best_test_path}"
         )
     
-        self.stats.end_training()
+    def _log_model_params(self):
-        self.stats.report(self.logger)
+        """输出模型可训练参数数量"""
-
-        self._finalize_training(best_model, best_test_model)
-
-        # 输出参数量
         try:
             total_params = sum(
                 p.numel() for p in self.model.parameters() if p.requires_grad
@@ -276,14 +332,20 @@ class Trainer:
 
     @staticmethod
     def test(model, args, data_loader, scaler, logger, path=None):
+        """对模型进行评估并输出性能指标"""
+        # 加载模型检查点（如果提供了路径）
         if path:
             checkpoint = torch.load(path)
             model.load_state_dict(checkpoint["state_dict"])
             model.to(args["basic"]["device"])
 
+        # 设置为评估模式
         model.eval()
+        
+        # 收集预测和真实标签
         y_pred, y_true = [], []
 
+        # 不计算梯度的情况下进行预测
         with torch.no_grad():
             for data, target in data_loader:
                 label = target[..., : args["output_dim"]]
@@ -291,12 +353,14 @@ class Trainer:
                 y_pred.append(output)
                 y_true.append(label)
 
+        # 合并所有批次的预测结果
         if args["real_value"]:
             y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
         else:
             y_pred = torch.cat(y_pred, dim=0)
         y_true = torch.cat(y_true, dim=0)
 
+        # 计算并记录每个时间步的指标
         for t in range(y_true.shape[1]):
             mae, rmse, mape = all_metrics(
                 y_pred[:, t, ...],
@@ -308,6 +372,7 @@ class Trainer:
                 f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
             )
 
+        # 计算并记录平均指标
         mae, rmse, mape = all_metrics(
             y_pred, y_true, args["mae_thresh"], args["mape_thresh"]
         )