opt trainer

train.py

@@ -6,14 +6,12 @@ import utils.initializer as init
 from dataloader.loader_selector import get_dataloader
 from trainer.trainer_selector import select_trainer
 
-import cProfile
-
 def read_config(config_path):
     with open(config_path, "r") as file:
         config = yaml.safe_load(file)
     
     # 全局配置
-    device = "cuda:1" # 指定设备为cuda:0
+    device = "cpu" # 指定设备为cuda:0
     seed = 2023 # 随机种子
     epochs = 120
 
@@ -67,8 +65,8 @@ def main(debug=False):
     model_list = ["iTransformer"]
     # 指定数据集
     # dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]
-    # dataset_list = ["AirQuality"]
-    dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"]
+    dataset_list = ["AirQuality"]
+    # dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"]
 
     # 我的调试开关，不做测试就填 str(False)
     # os.environ["TRY"] = str(False)
@@ -99,4 +97,4 @@ def main(debug=False):
 
 if __name__ == "__main__":
     # 调试用
-    main(debug = False)
+    main(debug = True)

trainer/Trainer.py

@@ -10,47 +10,45 @@ from tqdm import tqdm
 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.args = args["train"]
+        
         # 模型和训练相关组件
-        self.model = model
-        self.loss = loss
-        self.optimizer = optimizer
-        self.lr_scheduler = lr_scheduler
+        self.model, self.loss, self.optimizer, self.lr_scheduler = model, loss, optimizer, lr_scheduler
+
         # 数据加载器
-        self.train_loader = train_loader
-        self.val_loader = val_loader
-        self.test_loader = test_loader
+        self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader
+        
         # 数据处理工具
         self.scaler = scaler
-        self.args = train_args
+        
         # 初始化路径、日志和统计
-        self._initialize_paths(train_args)
-        self._initialize_logger(train_args)
+        self._initialize_paths(self.args)
+        self._initialize_logger(self.args)
     
     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")
+        log_dir = args["log_dir"]
+        self.best_path = os.path.join(log_dir, "best_model.pth")
+        self.best_test_path = os.path.join(log_dir, "best_test_model.pth")
+        self.loss_figure_path = os.path.join(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']}")
+        log_dir = args["log_dir"]
+        if not args["debug"]:
+            os.makedirs(log_dir, exist_ok=True)
+        self.logger = get_logger(log_dir, name=self.model.__class__.__name__, debug=args["debug"])
+        self.logger.info(f"Experiment log path in: {log_dir}")
 
     def _run_epoch(self, epoch, dataloader, mode):
         """运行一个训练/验证/测试epoch"""
         # 设置模型模式和是否进行优化
-        if mode == "train": self.model.train(); optimizer_step = True
-        else: self.model.eval(); optimizer_step = False
+        self.model.train() if mode == "train" else self.model.eval()
+        optimizer_step = mode == "train"
 
         # 初始化变量
         total_loss = 0
@@ -60,105 +58,111 @@ class Trainer:
         # 训练/验证循环
         with torch.set_grad_enabled(optimizer_step):
             progress_bar = tqdm(
-                enumerate(dataloader), 
+                dataloader, 
                 total=len(dataloader), 
                 desc=f"{mode.capitalize()} Epoch {epoch}"
             )
-            for _, (data, target) in progress_bar:
-                # 转移数据
-                data = data.to(self.device)
-                target = target.to(self.device)
+            for data, target in progress_bar:
+                # 转移数据并提取标签
+                data, target = data.to(self.device), target.to(self.device)
                 label = target[..., : self.args["output_dim"]]
-                # 计算loss和反归一化loss
+                
+                # 计算输出
                 output = self.model(data)
+                
                 # 我的调试开关
                 if os.environ.get("TRY") == "True":
-                    print(f"[{'✅' if output.shape == label.shape else '❌'}]: output: {output.shape}, label: {label.shape}")
+                    status = '✅' if output.shape == label.shape else '❌'
+                    print(f"[{status}]: output: {output.shape}, label: {label.shape}")
                     assert False
+                
+                # 计算损失
                 loss = self.loss(output, label)
                 d_output = self.scaler.inverse_transform(output)
                 d_label = self.scaler.inverse_transform(label)
                 d_loss = self.loss(d_output, d_label)
+                
                 # 累积损失和预测结果
                 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)
+        y_pred, y_true = torch.cat(y_pred, dim=0), 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} "
             f"MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s"
         )
+        
         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")
+        best_model = best_test_model = None
+        best_loss = best_test_loss = float("inf")
         not_improved_count = 0
+        
         # 开始训练
         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)
+            train_epoch_loss = self._run_epoch(epoch, self.train_loader, "train")
+            val_epoch_loss = self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
+            test_epoch_loss = self._run_epoch(epoch, self.test_loader, "test")
+            
             # 检查梯度爆炸
             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_loss, not_improved_count = val_epoch_loss, 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._finalize_training(best_model, best_test_model)
     
     def _should_early_stop(self, not_improved_count):
         """检查是否满足早停条件"""
-        if (
-            self.args["early_stop"]
-            and not_improved_count == self.args["early_stop_patience"]
-        ):
+        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."
             )
@@ -190,58 +194,43 @@ class Trainer:
     @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:
+        # 验证参数类型
+        if not isinstance(args, dict):
             raise ValueError(f"Unsupported args type: {type(args)}")
         
+        # 确定设备和输出维度
+        is_full_config = "basic" in args
+        device = args["basic"]["device"] if is_full_config else next(model.parameters()).device
+        output_dim = args["train"]["output_dim"] if is_full_config else args["output_dim"]
+        
+        # 获取metrics参数
+        train_args = args["train"] if is_full_config else args
+        mae_thresh, mape_thresh = train_args["mae_thresh"], train_args["mape_thresh"]
+        
         # 加载模型检查点（如果提供了路径）
         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)
-                
+                data, target = data.to(device), 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]):
@@ -254,9 +243,5 @@ class Trainer:
             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))
+        avg_mae, avg_rmse, avg_mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh)
+        logger.info(f"Average Horizon, MAE: {avg_mae:.4f}, RMSE: {avg_rmse:.4f}, MAPE: {avg_mape:.4f}")