config/args_parser.py

@@ -16,38 +16,4 @@ def parse_args():
     else:
         raise ValueError("Configuration file path must be provided using --config")
         
-    # Update configuration with command-line arguments
-    # Merge 'basic' configuration into the root dictionary
-    # config.update(config.get('basic', {}))
-
-    # Add adaptive configuration based on external commands
-    if "data" in config and "type" in config["data"]:
-        config["data"]["type"] = config["basic"].get("dataset", config["data"]["type"])
-    if "model" in config and "type" in config["model"]:
-        config["model"]["type"] = config["basic"].get("model", config["model"]["type"])
-    if "model" in config and "rnn_units" in config["model"]:
-        config["model"]["rnn_units"] = config["basic"].get(
-            "rnn", config["model"]["rnn_units"]
-        )
-    if "model" in config and "embed_dim" in config["model"]:
-        config["model"]["embed_dim"] = config["basic"].get(
-            "emb", config["model"]["embed_dim"]
-        )
-    if "data" in config and "sample" in config["data"]:
-        config["data"]["sample"] = config["basic"].get(
-            "sample", config["data"]["sample"]
-        )
-    if "train" in config and "device" in config["train"]:
-        config["train"]["device"] = config["basic"].get(
-            "device", config["train"]["device"]
-        )
-    if "train" in config and "debug" in config["train"]:
-        config["train"]["debug"] = config["basic"].get(
-            "debug", config["train"]["debug"]
-        )
-    if "cuda" in config:
-        config["cuda"] = config["basic"].get("cuda", config["cuda"])
-    if "mode" in config:
-        config["mode"] = config["basic"].get("mode", config["mode"])
-
     return config

generate_launch_configs.py

@@ -1,134 +0,0 @@
-import os
-import re
-
-# 配置路径
-CONFIG_DIR = "/user/czzhangheng/code/TrafficWheel/config"
-LAUNCH_FILE = "/user/czzhangheng/code/TrafficWheel/.vscode/launch.json"
-
-# 遍历所有yaml文件
-def find_all_yaml_files(directory):
-    yaml_files = []
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            if file.endswith(".yaml") and not file.startswith("BJTaxi"):
-                yaml_files.append(os.path.join(root, file))
-    return yaml_files
-
-# 生成launch配置字符串
-def generate_launch_config_string(yaml_files):
-    config_strings = []
-    
-    for file_path in yaml_files:
-        # 提取模型名和数据集名
-        relative_path = os.path.relpath(file_path, CONFIG_DIR)
-        model_name = relative_path.split(os.sep)[0]
-        dataset_name = os.path.splitext(os.path.basename(file_path))[0]
-        
-        # 处理v2版本
-        if "v2_" in dataset_name:
-            model_display_name = f"{model_name}_v2"
-            dataset_display_name = dataset_name.replace("v2_", "")
-        else:
-            model_display_name = model_name
-            dataset_display_name = dataset_name
-        
-        # 生成配置字符串
-        config_string = f'''
-        {{
-            "name": "{model_display_name}: {dataset_display_name}",
-            "type": "debugpy",
-            "request": "launch",
-            "program": "run.py",
-            "console": "integratedTerminal",
-            "args": "--config ./config/{model_name}/{os.path.basename(file_path)}"
-        }}'''
-        
-        config_strings.append(config_string)
-    
-    return ",".join(config_strings)
-
-# 读取现有的launch.json文件，提取配置名称
-def get_existing_config_names():
-    with open(LAUNCH_FILE, 'r') as f:
-        content = f.read()
-    
-    # 提取所有配置名称
-    name_pattern = re.compile(r'"name"\s*:\s*"([^"]+)"')
-    matches = name_pattern.findall(content)
-    
-    return set(matches)
-
-# 生成新的配置，过滤掉已存在的
-def generate_new_configs(yaml_files, existing_names):
-    new_configs = []
-    
-    for file_path in yaml_files:
-        # 提取模型名和数据集名
-        relative_path = os.path.relpath(file_path, CONFIG_DIR)
-        model_name = relative_path.split(os.sep)[0]
-        dataset_name = os.path.splitext(os.path.basename(file_path))[0]
-        
-        # 处理v2版本
-        if "v2_" in dataset_name:
-            model_display_name = f"{model_name}_v2"
-            dataset_display_name = dataset_name.replace("v2_", "")
-        else:
-            model_display_name = model_name
-            dataset_display_name = dataset_name
-        
-        # 生成配置名称
-        config_name = f"{model_display_name}: {dataset_display_name}"
-        
-        # 如果配置不存在，则添加
-        if config_name not in existing_names:
-            new_configs.append(file_path)
-    
-    return new_configs
-
-# 更新launch.json文件
-def update_launch_json(new_configs_string):
-    with open(LAUNCH_FILE, 'r') as f:
-        content = f.read()
-    
-    # 找到configurations数组的结束位置
-    configs_end_match = re.search(r'\s*\]\s*\}', content)
-    if not configs_end_match:
-        return False
-    
-    # 插入新的配置
-    insert_pos = configs_end_match.start()
-    new_content = content[:insert_pos] + new_configs_string + content[insert_pos:]
-    
-    # 保存文件
-    with open(LAUNCH_FILE, 'w') as f:
-        f.write(new_content)
-    
-    return True
-
-# 主函数
-def main():
-    # 查找所有yaml文件
-    yaml_files = find_all_yaml_files(CONFIG_DIR)
-    
-    # 获取现有配置名称
-    existing_names = get_existing_config_names()
-    
-    # 生成新的配置，过滤掉已存在的
-    new_config_files = generate_new_configs(yaml_files, existing_names)
-    
-    if not new_config_files:
-        print("No new configurations to add")
-        return
-    
-    # 生成新的配置字符串
-    new_configs_string = generate_launch_config_string(new_config_files)
-    
-    # 更新launch.json文件
-    if update_launch_json(new_configs_string):
-        print(f"Added {len(new_config_files)} new launch configurations")
-        print(f"Total configurations: {len(existing_names) + len(new_config_files)}")
-    else:
-        print("Failed to update launch.json")
-
-if __name__ == "__main__":
-    main()

mypy.ini

@@ -1,4 +0,0 @@
-[mypy]
-explicit_package_bases = True
-ignore_missing_imports = True
-no_site_packages = True

run_tests.sh

@@ -1,95 +0,0 @@
-#!/bin/bash
-
-# 设置默认模型名和数据集列表
-MODEL_NAME="STAEFormer"
-DATASETS=(
-    "METR-LA"
-    "PEMS-BAY"
-    "NYCBike-InFlow"
-    "NYCBike-OutFlow"
-    "AirQuality"
-    "SolarEnergy"
-)
-
-# 初始化统计变量
-success_count=0
-failure_count=0
-missing_count=0
-total_count=0
-success_datasets=()
-failure_datasets=()
-missing_datasets=()
-
-# 检查是否有参数传入来覆盖默认值
-if [ $# -gt 0 ]; then
-    MODEL_NAME=$1
-    # 如果传入了更多参数，使用它们作为数据集列表
-    if [ $# -gt 1 ]; then
-        DATASETS=(${@:2})
-    fi
-fi
-
-echo "使用模型: $MODEL_NAME"
-echo "数据集列表: ${DATASETS[*]}"
-echo "开始测试..."
-echo ""
-
-# 循环测试每个数据集
-for dataset in "${DATASETS[@]}"; do
-    total_count=$((total_count + 1))
-    # 构建配置文件路径
-    CONFIG_PATH="config/${MODEL_NAME}/${dataset}.yaml"
-    
-    echo "测试数据集: $dataset"
-    echo "使用配置文件: $CONFIG_PATH"
-    
-    # 检查配置文件是否存在
-    if [ ! -f "$CONFIG_PATH" ]; then
-        echo "错误: 配置文件 $CONFIG_PATH 不存在!"
-        missing_count=$((missing_count + 1))
-        missing_datasets+=("$dataset")
-        echo "----------------------------------------"
-        continue
-    fi
-    
-    # 执行测试命令，同时捕获输出并显示在控制台上
-    echo "执行: python run.py --config $CONFIG_PATH"
-    output=$(python run.py --config "$CONFIG_PATH" 2>&1 | tee /dev/tty)
-
-    # 如果没有找到明确的标记，回退到检查退出码
-    if [ $? -eq 0 ]; then
-        echo "数据集 $dataset 测试成功! (基于退出码)"
-        success_count=$((success_count + 1))
-        success_datasets+=("$dataset")
-    else
-        echo "数据集 $dataset 测试失败! (基于退出码)"
-        failure_count=$((failure_count + 1))
-        failure_datasets+=("$dataset")
-    fi
-    
-    echo "----------------------------------------"
-done
-
-# 输出总结
-echo "======================================="
-echo "测试总结"
-echo "======================================="
-echo "总数据集数量: $total_count"
-echo "成功数量: $success_count"
-echo "失败数量: $failure_count"
-echo "缺失配置文件数量: $missing_count"
-
-if [ ${#success_datasets[@]} -gt 0 ]; then
-    echo "成功的数据集: ${success_datasets[*]}"
-fi
-
-if [ ${#failure_datasets[@]} -gt 0 ]; then
-    echo "失败的数据集: ${failure_datasets[*]}"
-fi
-
-if [ ${#missing_datasets[@]} -gt 0 ]; then
-    echo "缺失配置的数据集: ${missing_datasets[*]}"
-fi
-
-echo "======================================="
-echo "所有测试完成!"

train.py

@@ -0,0 +1,63 @@
+import yaml
+import torch
+
+import utils.initializer as init
+from dataloader.loader_selector import get_dataloader
+from trainer.trainer_selector import select_trainer
+
+def run(config):
+    init.init_seed(config["basic"]["seed"])
+    model = init.init_model(config)
+    train_loader, val_loader, test_loader, scaler, *extra_data = get_dataloader(
+        config, normalizer=config["data"]["normalizer"], single=False
+    )
+    loss = init.init_loss(config, scaler)
+    optimizer, lr_scheduler = init.init_optimizer(model, config["train"])
+    init.create_logs(config)
+    trainer = select_trainer(
+        model,
+        loss, optimizer,
+        train_loader, val_loader, test_loader, scaler,
+        config,
+        lr_scheduler, extra_data,
+    )
+
+    # 开始训练
+    match config["basic"]["mode"]:
+        case "train":
+            trainer.train()
+        case "test":
+            model.load_state_dict(
+                torch.load(
+                    f"./pre-trained/{config['basic']['model']}/{config['basic']['dataset']}.pth",
+                    map_location=config["basic"]["device"],
+                    weights_only=True,
+                )
+            )
+            trainer.test(
+                model.to(config["basic"]["device"]),
+                trainer.args, test_loader, scaler,
+                trainer.logger,
+            )
+        case _:
+            raise ValueError(f"Unsupported mode: {config['basic']['mode']}")
+    
+
+if __name__ == "__main__":
+    # 指定模型
+    model_list = ["HI"] 
+    # 指定数据集
+    dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]
+    device = "cuda:0" # 指定设备
+    seed = 2023 # 随机种子
+    for model in model_list:
+        for dataset in dataset_list:
+            config_path = f"./config/{model}/{dataset}.yaml"
+            with open(config_path, "r") as file:
+                config = yaml.safe_load(file)
+            config["basic"]["device"] = device
+            config["basic"]["seed"] = seed
+            print(f"\nRunning {model} on {dataset} with seed {seed} on {device}")
+            print(f"config: {config}")
+            run(config)
+

trainer/Trainer.py

@@ -8,125 +8,31 @@ from utils.logger import get_logger
 from utils.loss_function import all_metrics
 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 = []
-        self.infer_time_list = []
-        self.total_iters = 0
-        self.start_time = None
-        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):
-        """记录单步耗时和总迭代次数"""
-        if mode == "train":
-            self.train_time_list.append(duration)
-        else:
-            self.infer_time_list.append(duration)
-        self.total_iters += 1
-
-    def record_memory_usage(self):
-        """记录当前 GPU 和 CPU 内存占用"""
-        process = psutil.Process(os.getpid())
-        cpu_mem = process.memory_info().rss / (1024**2)
-
-        if torch.cuda.is_available():
-            gpu_mem = torch.cuda.max_memory_allocated(device=self.device) / (1024**2)
-            torch.cuda.reset_peak_memory_stats(device=self.device)
-        else:
-            gpu_mem = 0.0
-
-        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:
-            logger.warning("TrainingStats: start/end time not recorded properly.")
-            return
-
-        total_time = self.end_time - self.start_time
-        avg_gpu_mem = self._calculate_average(self.gpu_mem_usage_list)
-        avg_cpu_mem = self._calculate_average(self.cpu_mem_usage_list)
-        avg_train_time = self._calculate_average(self.train_time_list)
-        avg_infer_time = self._calculate_average(self.infer_time_list)
-        iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
-
-        logger.info("===== Training Summary =====")
-        logger.info(f"Total training time: {total_time:.2f} s")
-        logger.info(f"Total iterations: {self.total_iters}")
-        logger.info(f"Average iterations per second: {iters_per_sec:.2f}")
-        logger.info(f"Average GPU Memory Usage: {avg_gpu_mem:.2f} MB")
-        logger.info(f"Average CPU Memory Usage: {avg_cpu_mem:.2f} MB")
-        if avg_train_time:
-            logger.info(f"Average training step time: {avg_train_time * 1000:.2f} ms")
-        if avg_infer_time:
-            logger.info(f"Average inference step time: {avg_infer_time * 1000:.2f} ms")
-
-
 class Trainer:
     """模型训练器，负责整个训练流程的管理"""
     
-    def __init__(
-        self,
-        model,
-        loss,
-        optimizer,
-        train_loader,
-        val_loader,
-        test_loader,
-        scaler,
-        args,
-        lr_scheduler=None,
-    ):
+    def __init__(self, model, loss, optimizer,
+                train_loader, val_loader, test_loader, scaler,
+                args, lr_scheduler=None,):
         # 设备和基本参数
+        self.config = args
         self.device = args["basic"]["device"]
         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 = train_args
-        
-        # 统计信息
-        self.train_per_epoch = len(train_loader)
-        self.val_per_epoch = len(val_loader) if val_loader else 0
-
         # 初始化路径、日志和统计
         self._initialize_paths(train_args)
         self._initialize_logger(train_args)
-        self._initialize_stats()
     
     def _initialize_paths(self, args):
         """初始化模型保存路径"""
@@ -138,24 +44,14 @@ class Trainer:
         """初始化日志记录器"""
         if not os.path.isdir(args["log_dir"]) and not args["debug"]:
             os.makedirs(args["log_dir"], exist_ok=True)
-        self.logger = get_logger(
-            args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
-        )
+        self.logger = get_logger(args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"])
         self.logger.info(f"Experiment log path in: {args['log_dir']}")
 
-    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
-        else:
-            self.model.eval()
-            optimizer_step = False
+        if mode == "train": self.model.train(); optimizer_step = True
+        else: self.model.eval(); optimizer_step = False
 
         # 初始化变量
         total_loss = 0
@@ -169,73 +65,42 @@ class Trainer:
                 total=len(dataloader), 
                 desc=f"{mode.capitalize()} Epoch {epoch}"
             )
-            
             for _, (data, target) in progress_bar:
-                # 记录步骤开始时间
-                start_time = time.time()
-
-                # 前向传播
+                # 转移数据
+                data = data.to(self.device)
+                target = target.to(self.device)
                 label = target[..., : self.args["output_dim"]]
-                output = self.model(data).to(self.device)
-                # if output.shape != label.shape:
-                #     import sys
-                #     print(f"[Wrong]: Output shape: {output.shape}, Label shape: {label.shape}")
-                #     sys.exit(1)
-                # else:
-                #     import sys
-                #     print(f"[Right]: Output shape: {output.shape}, Label shape: {label.shape}")
-                #     sys.exit(0)
+                # 计算loss和反归一化loss
+                output = self.model(data)
                 loss = self.loss(output, label)
-
-                # 反归一化
                 d_output = self.scaler.inverse_transform(output)
                 d_label = 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()
-                
-                # 反归一化的loss
                 d_loss = self.loss(d_output, d_label)
-
-                # 记录步骤时间和内存使用
-                step_time = time.time() - start_time
-                self.stats.record_step_time(step_time, mode)
-
                 # 累积损失和预测结果
                 total_loss += d_loss.item()
                 y_pred.append(d_output.detach().cpu())
                 y_true.append(d_label.detach().cpu())
-
+                # 反向传播和优化（仅在训练模式）
+                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()
                 # 更新进度条
                 progress_bar.set_postfix(loss=d_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"]
-        )
+        mae, rmse, mape = all_metrics(y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"])
         self.logger.info(
-            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"
+            f"Epoch #{epoch:02d}: {mode.capitalize():<5} "
+            f"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
 
     def train_epoch(self, epoch):
@@ -248,28 +113,22 @@ 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
@@ -278,30 +137,19 @@ class Trainer:
                 self.logger.info("Best validation model saved!")
             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 (
@@ -331,20 +179,35 @@ class Trainer:
     def _finalize_training(self, best_model, best_test_model):
         self.model.load_state_dict(best_model)
         self.logger.info("Testing on best validation model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
-
+        self.test(self.model, self.config, self.test_loader, self.scaler, self.logger)
         self.model.load_state_dict(best_test_model)
         self.logger.info("Testing on best test model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
+        self.test(self.model, self.config, self.test_loader, self.scaler, self.logger)
 
     @staticmethod
     def test(model, args, data_loader, scaler, logger, path=None):
         """对模型进行评估并输出性能指标"""
+        # 确定设备信息
+        device = None
+        output_dim = None
+        # 处理不同的参数格式
+        if isinstance(args, dict):
+            if "basic" in args:
+                # 完整配置情况
+                device = args["basic"]["device"]
+                output_dim = args["train"]["output_dim"]
+            else:
+                # 只有train_args情况，从模型获取设备
+                device = next(model.parameters()).device
+                output_dim = args["output_dim"]
+        else:
+            raise ValueError(f"Unsupported args type: {type(args)}")
+        
         # 加载模型检查点（如果提供了路径）
         if path:
             checkpoint = torch.load(path)
             model.load_state_dict(checkpoint["state_dict"])
-            model.to(args["basic"]["device"])
+            model.to(device)
 
         # 设置为评估模式
         model.eval()
@@ -355,27 +218,40 @@ class Trainer:
         # 不计算梯度的情况下进行预测
         with torch.no_grad():
             for data, target in data_loader:
-                label = target[..., : args["output_dim"]]
+                # 将数据和标签移动到指定设备
+                data = data.to(device)
+                target = target.to(device)
+                
+                label = target[..., : output_dim]
                 output = model(data)
                 y_pred.append(output.detach().cpu())
                 y_true.append(label.detach().cpu())
 
-        
         d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
         d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
 
+        # 获取metrics参数
+        if "basic" in args:
+            # 完整配置情况
+            mae_thresh = args["train"]["mae_thresh"]
+            mape_thresh = args["train"]["mape_thresh"]
+        else:
+            # 只有train_args情况
+            mae_thresh = args["mae_thresh"]
+            mape_thresh = args["mape_thresh"]
+        
         # 计算并记录每个时间步的指标
         for t in range(d_y_true.shape[1]):
             mae, rmse, mape = all_metrics(
                 d_y_pred[:, t, ...],
                 d_y_true[:, t, ...],
-                args["mae_thresh"],
-                args["mape_thresh"],
+                mae_thresh,
+                mape_thresh,
             )
             logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
 
         # 计算并记录平均指标
-        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh)
         logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
 
     @staticmethod

trainer/Trainer_bk.py

@@ -0,0 +1,420 @@
+import math
+import os
+import time
+import copy
+import psutil
+import torch
+from utils.logger import get_logger
+from utils.loss_function import all_metrics
+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 = []
+#         self.infer_time_list = []
+#         self.total_iters = 0
+#         self.start_time = None
+#         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):
+#         """记录单步耗时和总迭代次数"""
+#         if mode == "train":
+#             self.train_time_list.append(duration)
+#         else:
+#             self.infer_time_list.append(duration)
+#         self.total_iters += 1
+
+#     def record_memory_usage(self):
+#         """记录当前 GPU 和 CPU 内存占用"""
+#         process = psutil.Process(os.getpid())
+#         cpu_mem = process.memory_info().rss / (1024**2)
+
+#         if torch.cuda.is_available():
+#             gpu_mem = torch.cuda.max_memory_allocated(device=self.device) / (1024**2)
+#             torch.cuda.reset_peak_memory_stats(device=self.device)
+#         else:
+#             gpu_mem = 0.0
+
+#         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:
+#             logger.warning("TrainingStats: start/end time not recorded properly.")
+#             return
+
+#         total_time = self.end_time - self.start_time
+#         avg_gpu_mem = self._calculate_average(self.gpu_mem_usage_list)
+#         avg_cpu_mem = self._calculate_average(self.cpu_mem_usage_list)
+#         avg_train_time = self._calculate_average(self.train_time_list)
+#         avg_infer_time = self._calculate_average(self.infer_time_list)
+#         iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
+
+#         logger.info("===== Training Summary =====")
+#         logger.info(f"Total training time: {total_time:.2f} s")
+#         logger.info(f"Total iterations: {self.total_iters}")
+#         logger.info(f"Average iterations per second: {iters_per_sec:.2f}")
+#         logger.info(f"Average GPU Memory Usage: {avg_gpu_mem:.2f} MB")
+#         logger.info(f"Average CPU Memory Usage: {avg_cpu_mem:.2f} MB")
+#         if avg_train_time:
+#             logger.info(f"Average training step time: {avg_train_time * 1000:.2f} ms")
+#         if avg_infer_time:
+#             logger.info(f"Average inference step time: {avg_infer_time * 1000:.2f} ms")
+
+
+class Trainer:
+    """模型训练器，负责整个训练流程的管理"""
+    
+    def __init__(
+        self,
+        model,
+        loss,
+        optimizer,
+        train_loader,
+        val_loader,
+        test_loader,
+        scaler,
+        args,
+        lr_scheduler=None,
+    ):
+        # 设备和基本参数
+        self.device = args["basic"]["device"]
+        self.config = args  # 保存完整的配置参数
+        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 = train_args
+        
+        # 统计信息
+        # self.train_per_epoch = len(train_loader)
+        # self.val_per_epoch = len(val_loader) if val_loader else 0
+
+        # 初始化路径、日志和统计
+        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")
+    
+    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(
+            args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
+        )
+        self.logger.info(f"Experiment log path in: {args['log_dir']}")
+    
+    # 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
+        else:
+            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}"
+            )
+            
+            for _, (data, target) in progress_bar:
+                # 记录步骤开始时间
+                start_time = time.time()
+
+                # 将数据和标签移动到指定设备
+                data = data.to(self.device)
+                target = target.to(self.device)
+                
+                # 前向传播
+                label = target[..., : self.args["output_dim"]]
+                output = self.model(data)
+                # if output.shape != label.shape:
+                #     import sys
+                #     print(f"[Wrong]: Output shape: {output.shape}, Label shape: {label.shape}")
+                #     sys.exit(1)
+                # else:
+                #     import sys
+                #     print(f"[Right]: Output shape: {output.shape}, Label shape: {label.shape}")
+                #     sys.exit(0)
+                loss = self.loss(output, label)
+
+                # 反归一化
+                d_output = self.scaler.inverse_transform(output)
+                d_label = 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()
+                
+                # 反归一化的loss
+                d_loss = self.loss(d_output, d_label)
+
+                # 记录步骤时间和内存使用
+                # step_time = time.time() - start_time
+                # self.stats.record_step_time(step_time, mode)
+
+                # 累积损失和预测结果
+                total_loss += d_loss.item()
+                y_pred.append(d_output.detach().cpu())
+                y_true.append(d_label.detach().cpu())
+
+                # 更新进度条
+                progress_bar.set_postfix(loss=d_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"]
+        )
+        self.logger.info(
+            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
+
+    def train_epoch(self, epoch):
+        return self._run_epoch(epoch, self.train_loader, "train")
+
+    def val_epoch(self, epoch):
+        return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
+
+    def test_epoch(self, epoch):
+        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
+                best_model = copy.deepcopy(self.model.state_dict())
+                self.logger.info("Best validation model saved!")
+            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"]
+        ):
+            self.logger.info(
+                f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops."
+            )
+            return True
+        return False
+    
+    def _save_best_models(self, best_model, best_test_model):
+        """保存最佳模型到文件"""
+        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}"
+        )
+    
+    def _log_model_params(self):
+        """输出模型可训练参数数量"""
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
+        self.logger.info(f"Trainable params: {total_params}")
+        
+
+    def _finalize_training(self, best_model, best_test_model):
+        self.model.load_state_dict(best_model)
+        self.logger.info("Testing on best validation model")
+        self.test(self.model, self.config, self.test_loader, self.scaler, self.logger)
+
+        self.model.load_state_dict(best_test_model)
+        self.logger.info("Testing on best test model")
+        self.test(self.model, self.config, self.test_loader, self.scaler, self.logger)
+
+    @staticmethod
+    def test(model, args, data_loader, scaler, logger, path=None):
+        """对模型进行评估并输出性能指标"""
+        # 确定设备信息
+        device = None
+        output_dim = None
+        
+        # 处理不同的参数格式
+        if isinstance(args, dict):
+            if "basic" in args:
+                # 完整配置情况
+                device = args["basic"]["device"]
+                output_dim = args["train"]["output_dim"]
+            else:
+                # 只有train_args情况
+                # 从模型获取设备
+                device = next(model.parameters()).device
+                output_dim = args["output_dim"]
+        else:
+            raise ValueError(f"Unsupported args type: {type(args)}")
+        
+        # 加载模型检查点（如果提供了路径）
+        if path:
+            checkpoint = torch.load(path)
+            model.load_state_dict(checkpoint["state_dict"])
+            model.to(device)
+
+        # 设置为评估模式
+        model.eval()
+        
+        # 收集预测和真实标签
+        y_pred, y_true = [], []
+
+        # 不计算梯度的情况下进行预测
+        with torch.no_grad():
+            for data, target in data_loader:
+                # 将数据和标签移动到指定设备
+                data = data.to(device)
+                target = target.to(device)
+                
+                label = target[..., : output_dim]
+                output = model(data)
+                y_pred.append(output.detach().cpu())
+                y_true.append(label.detach().cpu())
+
+        
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
+
+        # 获取metrics参数
+        if "basic" in args:
+            # 完整配置情况
+            mae_thresh = args["train"]["mae_thresh"]
+            mape_thresh = args["train"]["mape_thresh"]
+        else:
+            # 只有train_args情况
+            mae_thresh = args["mae_thresh"]
+            mape_thresh = args["mape_thresh"]
+        
+        # 计算并记录每个时间步的指标
+        for t in range(d_y_true.shape[1]):
+            mae, rmse, mape = all_metrics(
+                d_y_pred[:, t, ...],
+                d_y_true[:, t, ...],
+                mae_thresh,
+                mape_thresh,
+            )
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
+
+        # 计算并记录平均指标
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh)
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
+
+    @staticmethod
+    def _compute_sampling_threshold(global_step, k):
+        return k / (k + math.exp(global_step / k))

trainer/Trainer_old.py

@@ -1,229 +0,0 @@
-import math
-import os
-import time
-import copy
-from tqdm import tqdm
-
-import torch
-from utils.logger import get_logger
-from utils.loss_function import all_metrics
-from utils.training_stats import TrainingStats
-
-
-class Trainer:
-    def __init__(
-        self,
-        model,
-        loss,
-        optimizer,
-        train_loader,
-        val_loader,
-        test_loader,
-        scaler,
-        args,
-        lr_scheduler=None,
-    ):
-        self.model = model
-        self.loss = loss
-        self.optimizer = optimizer
-        self.train_loader = train_loader
-        self.val_loader = val_loader
-        self.test_loader = test_loader
-        self.scaler = scaler
-        self.args = 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.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
-        if not os.path.isdir(args["log_dir"]) and not args["debug"]:
-            os.makedirs(args["log_dir"], exist_ok=True)
-        self.logger = get_logger(
-            args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
-        )
-        self.logger.info(f"Experiment log path in: {args['log_dir']}")
-        # Stats tracker
-        self.stats = TrainingStats(device=args["device"])
-
-    def _run_epoch(self, epoch, dataloader, mode):
-        if mode == "train":
-            self.model.train()
-            optimizer_step = True
-        else:
-            self.model.eval()
-            optimizer_step = False
-
-        total_loss = 0
-        epoch_time = time.time()
-
-        with torch.set_grad_enabled(optimizer_step):
-            with tqdm(
-                total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
-            ) as pbar:
-                for batch_idx, (data, target) in enumerate(dataloader):
-                    start_time = time.time()
-                    label = target[..., : self.args["output_dim"]]
-                    output = self.model(data).to(self.args["device"])
-
-                    if self.args["real_value"]:
-                        output = self.scaler.inverse_transform(output)
-
-                    loss = self.loss(output, 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()
-
-                    if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0:
-                        self.logger.info(
-                            f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}"
-                        )
-
-                    # 更新 tqdm 的进度
-                    pbar.update(1)
-                    pbar.set_postfix(loss=loss.item())
-
-        avg_loss = total_loss / len(dataloader)
-        self.logger.info(
-            f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
-        )
-        # 记录内存
-        self.stats.record_memory_usage()
-        return avg_loss
-
-    def train_epoch(self, epoch):
-        return self._run_epoch(epoch, self.train_loader, "train")
-
-    def val_epoch(self, epoch):
-        return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
-
-    def test_epoch(self, epoch):
-        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):
-            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
-                best_model = copy.deepcopy(self.model.state_dict())
-                self.logger.info("Best validation model saved!")
-            else:
-                not_improved_count += 1
-
-            if (
-                self.args["early_stop"]
-                and not_improved_count == self.args["early_stop_patience"]
-            ):
-                self.logger.info(
-                    f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops."
-                )
-                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"]:
-            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()
-        self.stats.report(self.logger)
-        try:
-            total_params = sum(
-                p.numel() for p in self.model.parameters() if p.requires_grad
-            )
-            self.logger.info(f"Trainable params: {total_params}")
-        except Exception:
-            pass
-        self._finalize_training(best_model, best_test_model)
-
-    def _finalize_training(self, best_model, best_test_model):
-        self.model.load_state_dict(best_model)
-        self.logger.info("Testing on best validation model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
-
-        self.model.load_state_dict(best_test_model)
-        self.logger.info("Testing on best test model")
-        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
-
-    @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["device"])
-
-        model.eval()
-        y_pred, y_true = [], []
-
-        with torch.no_grad():
-            for data, target in data_loader:
-                label = target[..., : args["output_dim"]]
-                output = model(data)
-                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)
-
-        # 你在这里需要把y_pred和y_true保存下来
-        # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1]
-        # torch.save(y_true, "./test/PEMS08/y_true.pt") # [3566,12,170,1]
-
-        for t in range(y_true.shape[1]):
-            mae, rmse, mape = all_metrics(
-                y_pred[:, t, ...],
-                y_true[:, t, ...],
-                args["mae_thresh"],
-                args["mape_thresh"],
-            )
-            logger.info(
-                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"]
-        )
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
-
-    @staticmethod
-    def _compute_sampling_threshold(global_step, k):
-        return k / (k + math.exp(global_step / k))

utils/initializer.py

@@ -9,9 +9,9 @@ import os
 import yaml
 
 
-def init_model(args):
-    device = args["device"]
-    model = model_selector(args).to(device)
+def init_model(config):
+    device = config["basic"]["device"]
+    model = model_selector(config).to(device)
     for p in model.parameters():
         if p.dim() > 1:
             nn.init.xavier_uniform_(p)