trainer修改

mypy.ini

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

trainer/DCRNN_Trainer.py

@@ -2,6 +2,7 @@ import math
 import os
 import time
 import copy
+import psutil
 from tqdm import tqdm
 
 import torch
@@ -23,34 +24,56 @@ class Trainer:
         args,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
+        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 = 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(
             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 _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
@@ -58,54 +81,77 @@ 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):
-            with tqdm(
+            progress_bar = tqdm(
-                total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
+                enumerate(dataloader), 
-            ) as pbar:
+                total=len(dataloader), 
-                for batch_idx, (data, target) in enumerate(dataloader):
+                desc=f"{mode.capitalize()} Epoch {epoch}"
-                    start_time = time.time()
-                    label = target[..., : self.args["output_dim"]]
-                    output = self.model(data, labels=label.clone()).to(
-                        self.args["device"]
             )
             
-                    if self.args["real_value"]:
+            for _, (data, target) in progress_bar:
-                        output = self.scaler.inverse_transform(output)
+                # 记录步骤开始时间
-                        label = self.scaler.inverse_transform(label)
+                start_time = time.time()
 
+                # 前向传播
+                label = target[..., : self.args["output_dim"]]
+                output = self.model(data, labels=label.clone()).to(self.device)
                 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 += loss.item()
 
-                    if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0:
+                # 累积损失和预测结果
-                        self.logger.info(
+                total_loss += d_loss.item()
-                            f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}"
+                y_pred.append(d_output.detach().cpu())
-                        )
+                y_true.append(d_label.detach().cpu())
 
-                    # 更新 tqdm 的进度
+                # 更新进度条
-                    pbar.update(1)
+                progress_bar.set_postfix(loss=d_loss.item())
-                    pbar.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)
-        self.logger.info(
+        
-            f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
+        # 计算并记录指标
+        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):
@@ -118,21 +164,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
@@ -141,6 +195,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"]
@@ -148,30 +227,22 @@ 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}"
         )
     
-        # 输出统计与参数
+    def _log_model_params(self):
-        self.stats.end_training()
+        """输出模型可训练参数数量"""
-        self.stats.report(self.logger)
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
-        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)
@@ -184,44 +255,44 @@ 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["device"])
+            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"]]
-                output = model(data, labels=label.clone()).to(args["device"])
+                output = model(data, labels=label.clone())
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
 
-        y_pred = torch.cat(y_pred, dim=0)
+        # 反归一化
-        y_true = torch.cat(y_true, dim=0)
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        if args["real_value"]:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-            y_pred = scaler.inverse_transform(y_pred)
-            y_true = scaler.inverse_transform(y_true)
 
-        for t in range(y_true.shape[1]):
+        # 计算并记录每个时间步的指标
+        for t in range(d_y_true.shape[1]):
             mae, rmse, mape = all_metrics(
-                y_pred[:, t, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

trainer/E32Trainer.py

@@ -23,44 +23,65 @@ class Trainer:
         global_config,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
         self.device = global_config["basic"]["device"]
         train_config = global_config["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_config
-        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(train_config["log_dir"], "best_model.pth")
+        self._initialize_paths(train_config)
-        self.best_test_path = os.path.join(
+        self._initialize_logger(train_config)
-            train_config["log_dir"], "best_test_model.pth"
+        self._initialize_stats()
-        )
-        self.loss_figure_path = os.path.join(train_config["log_dir"], "loss.png")
     
-        # Initialize logger
+    def _initialize_paths(self, args):
-        if not os.path.isdir(train_config["log_dir"]) and not train_config["debug"]:
+        """初始化模型保存路径"""
-            os.makedirs(train_config["log_dir"], exist_ok=True)
+        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(
-            train_config["log_dir"],
+            args["log_dir"],
             name=self.model.__class__.__name__,
-            debug=train_config["debug"],
+            debug=args["debug"],
         )
-        self.logger.info(f"Experiment log path in: {train_config['log_dir']}")
+        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"""
+        # 设置模型模式和是否进行优化
         is_train = mode == "train"
         self.model.train() if is_train else self.model.eval()
+        
+        # 初始化变量
         total_loss = 0.0
         epoch_time = time.time()
+        y_pred, y_true = [], []
 
         with (
             torch.set_grad_enabled(is_train),
@@ -85,10 +106,12 @@ class Trainer:
 
                 # compute loss
                 label = target[..., : self.args["output_dim"]]
-                if self.args["real_value"]:
-                    output = self.scaler.inverse_transform(output)
                 loss = self.loss(output, label)
 
+                # 反归一化
+                d_output = self.scaler.inverse_transform(output)
+                d_label = self.scaler.inverse_transform(label)
+
                 # backward / step
                 if is_train:
                     loss.backward()
@@ -98,22 +121,39 @@ class Trainer:
                         )
                     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 += loss.item()
+                total_loss += d_loss.item()
+
+                # 累积预测结果
+                y_pred.append(d_output.detach().cpu())
+                y_true.append(d_label.detach().cpu())
 
                 # logging
                 if is_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}"
+                        f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)}  Loss: {d_loss.item():.6f}"
                     )
 
                 pbar.update(1)
-                pbar.set_postfix(loss=loss.item())
+                pbar.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"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
+            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()
@@ -129,21 +169,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
@@ -152,6 +200,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"]
@@ -159,30 +232,22 @@ 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}"
         )
     
-        # 输出统计与参数
+    def _log_model_params(self):
-        self.stats.end_training()
+        """输出模型可训练参数数量"""
-        self.stats.report(self.logger)
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
-        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)
@@ -195,51 +260,44 @@ class Trainer:
 
     @staticmethod
     def test(model, args, data_loader, scaler, logger, path=None):
-        global_config = args
+        """对模型进行评估并输出性能指标"""
-        device = global_config["basic"]["device"]
+        # 加载模型检查点（如果提供了路径）
-        args = global_config["train"]
         if path:
             checkpoint = torch.load(path)
             model.load_state_dict(checkpoint["state_dict"])
-            model.to(device)
+            model.to(args["basic"]["device"])
 
+        # 设置为评估模式
         model.eval()
+        
+        # 收集预测和真实标签
         y_pred, y_true = [], []
 
+        # 不计算梯度的情况下进行预测
         with torch.no_grad():
             for data, target, cycle_index in data_loader:
                 label = target[..., : args["output_dim"]]
                 output = model(data, cycle_index)
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
 
-        if args["real_value"]:
+        # 反归一化
-            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        else:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-            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]
+        for t in range(d_y_true.shape[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, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

trainer/EXP_trainer.py

@@ -2,6 +2,7 @@ import math
 import os
 import time
 import copy
+import psutil
 from tqdm import tqdm
 
 import torch
@@ -23,34 +24,56 @@ class Trainer:
         args,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
+        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 = 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(
             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 _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
@@ -58,52 +81,77 @@ 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):
-            with tqdm(
+            progress_bar = tqdm(
-                total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
+                enumerate(dataloader), 
-            ) as pbar:
+                total=len(dataloader), 
-                for batch_idx, (data, target) in enumerate(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.args["device"])
+                output = self.model(data).to(self.device)
-
-                    if self.args["real_value"]:
-                        output = self.scaler.inverse_transform(output)
-
                 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 += loss.item()
+                # 累积损失和预测结果
+                total_loss += d_loss.item()
+                y_pred.append(d_output.detach().cpu())
+                y_true.append(d_label.detach().cpu())
 
-                    if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0:
+                # 更新进度条
-                        self.logger.info(
+                progress_bar.set_postfix(loss=d_loss.item())
-                            f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}"
-                        )
 
-                    # 更新 tqdm 的进度
+        # 合并所有批次的预测结果
-                    pbar.update(1)
+        y_pred = torch.cat(y_pred, dim=0)
-                    pbar.set_postfix(loss=loss.item())
+        y_true = torch.cat(y_true, dim=0)
 
+        # 计算平均损失
         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"
+        # 计算并记录指标
+        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):
@@ -116,21 +164,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
@@ -139,6 +195,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"]
@@ -146,29 +227,22 @@ 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)
+        """输出模型可训练参数数量"""
-        try:
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
-            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)
@@ -181,48 +255,44 @@ 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["device"])
+            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"]]
                 output = model(data)
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
 
-        if args["real_value"]:
+        # 反归一化
-            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        else:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-            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]
+        for t in range(d_y_true.shape[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, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

trainer/PDG2SEQ_Trainer.py

@@ -23,35 +23,57 @@ class Trainer:
         args,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
+        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 = 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
         self.batches_seen = 0
         
-        # Paths for saving models and logs
+        # 统计信息
+        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")
     
-        # 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(
             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 _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
@@ -59,25 +81,40 @@ 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):
-            with tqdm(
+            progress_bar = tqdm(
-                total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
+                enumerate(dataloader), 
-            ) as pbar:
+                total=len(dataloader), 
-                for batch_idx, (data, target) in enumerate(dataloader):
+                desc=f"{mode.capitalize()} Epoch {epoch}"
+            )
+            
+            for batch_idx, (data, target) in progress_bar:
                 start_time = time.time()
                 self.batches_seen += 1
                 label = target[..., : self.args["output_dim"]].clone()
-                    output = self.model(data, target, self.batches_seen).to(
-                        self.args["device"]
-                    )
                 
-                    if self.args["real_value"]:
+                # 前向传播
-                        output = self.scaler.inverse_transform(output)
+                if mode == "train":
+                    output = self.model(data, target, self.batches_seen).to(self.device)
+                else:
+                    output = self.model(data, target).to(self.device)
                 
+                # 计算原始loss
                 loss = self.loss(output, label)
+
+                # 反归一化
+                d_output = self.scaler.inverse_transform(output)
+                d_label = self.scaler.inverse_transform(label)
+
+                # 反归一化的loss
+                d_loss = self.loss(d_output, d_label)
+
+                # 反向传播和优化（仅在训练模式）
                 if optimizer_step and self.optimizer is not None:
                     self.optimizer.zero_grad()
                     loss.backward()
@@ -88,26 +125,42 @@ class Trainer:
                         )
                     self.optimizer.step()
 
-                    # record step time
+                # 记录步骤时间
                 step_time = time.time() - start_time
                 self.stats.record_step_time(step_time, mode)
-                    total_loss += loss.item()
+                total_loss += d_loss.item()
+
+                # 累积预测结果
+                y_pred.append(d_output.detach().cpu())
+                y_true.append(d_label.detach().cpu())
 
                 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}"
+                        f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
                     )
 
                 # 更新 tqdm 的进度
-                    pbar.update(1)
+                progress_bar.update(1)
-                    pbar.set_postfix(loss=loss.item())
+                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)
-        self.logger.info(
+        
-            f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
+        # 计算并记录指标
+        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):
@@ -120,21 +173,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
@@ -143,6 +204,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._log_model_params()
+        
+        # 最终评估
+        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"]
@@ -150,30 +236,22 @@ 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}"
         )
     
-        # 输出统计与参数
+    def _log_model_params(self):
-        self.stats.end_training()
+        """输出模型可训练参数数量"""
-        self.stats.report(self.logger)
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
-        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)
@@ -186,44 +264,44 @@ 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["device"])
+            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"]].clone()
                 output = model(data, target)
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
 
-        if args["real_value"]:
+        # 反归一化
-            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        else:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-            y_pred = torch.cat(y_pred, dim=0)
-        y_true = torch.cat(y_true, dim=0)
 
-        for t in range(y_true.shape[1]):
+        # 计算并记录每个时间步的指标
+        for t in range(d_y_true.shape[1]):
             mae, rmse, mape = all_metrics(
-                y_pred[:, t, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

trainer/STMLP_Trainer.py

@@ -26,42 +26,35 @@ class Trainer:
         args,
         lr_scheduler=None,
     ):
+        # 设备和基本参数
+        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 = args["train"]
+        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.best_path = os.path.join(self.args["log_dir"], "best_model.pth")
+        self._initialize_paths(args, train_args)
-        self.best_test_path = os.path.join(self.args["log_dir"], "best_test_model.pth")
+        self._initialize_logger(train_args)
-        self.loss_figure_path = os.path.join(self.args["log_dir"], "loss.png")
+        self._initialize_stats()
-        self.pretrain_dir = (
-            f"./pre-train/{args['model']['type']}/{args['data']['type']}"
-        )
-        self.pretrain_path = os.path.join(self.pretrain_dir, "best_model.pth")
-        self.pretrain_best_path = os.path.join(self.pretrain_dir, "best_test_model.pth")
-
-        # Initialize logger
-        if not os.path.isdir(self.args["log_dir"]) and not self.args["debug"]:
-            os.makedirs(self.args["log_dir"], exist_ok=True)
-        if not os.path.isdir(self.pretrain_dir) and not self.args["debug"]:
-            os.makedirs(self.pretrain_dir, exist_ok=True)
-        self.logger = get_logger(
-            self.args["log_dir"],
-            name=self.model.__class__.__name__,
-            debug=self.args["debug"],
-        )
-        self.logger.info(f"Experiment log path in: {self.args['log_dir']}")
-        # Stats tracker
-        self.stats = TrainingStats(device=args["device"])
         
+        # 教师-学生蒸馏相关
         if self.args["teacher_stu"]:
             self.tmodel = self.loadTeacher(args)
         else:
@@ -71,8 +64,40 @@ class Trainer:
                 f"然后在config中配置train.teacher_stu模式为True开启蒸馏模式"
             )
     
+    def _initialize_paths(self, args, train_args):
+        """初始化模型保存路径"""
+        self.best_path = os.path.join(train_args["log_dir"], "best_model.pth")
+        self.best_test_path = os.path.join(train_args["log_dir"], "best_test_model.pth")
+        self.loss_figure_path = os.path.join(train_args["log_dir"], "loss.png")
+        self.pretrain_dir = (
+            f"./pre-train/{args['model']['type']}/{args['data']['type']}"
+        )
+        self.pretrain_path = os.path.join(self.pretrain_dir, "best_model.pth")
+        self.pretrain_best_path = os.path.join(self.pretrain_dir, "best_test_model.pth")
+        
+        # 创建预训练目录
+        if not os.path.isdir(self.pretrain_dir) and not train_args["debug"]:
+            os.makedirs(self.pretrain_dir, exist_ok=True)
+    
+    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"""
         # self.tmodel.eval()
+        # 设置模型模式和是否进行优化
         if mode == "train":
             self.model.train()
             optimizer_step = True
@@ -80,8 +105,10 @@ 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):
             with tqdm(
@@ -89,15 +116,17 @@ class Trainer:
             ) as pbar:
                 for batch_idx, (data, target) in enumerate(dataloader):
                     start_time = time.time()
-                    if self.args["teacher_stu"]:
                     label = target[..., : self.args["output_dim"]]
+                    
+                    if self.args["teacher_stu"]:
+                        # 教师-学生蒸馏模式
                         output, out_, _ = self.model(data)
                         gout, tout, sout = self.tmodel(data)
                         
-                        if self.args["real_value"]:
+                        # 计算原始loss
-                            output = self.scaler.inverse_transform(output)
-
                         loss1 = self.loss(output, label)
+                        
+                        # 计算蒸馏相关loss
                         scl = self.loss_cls(out_, sout)
                         kl_loss = nn.KLDivLoss(
                             reduction="batchmean", log_target=True
@@ -105,17 +134,22 @@ class Trainer:
                         gout = F.log_softmax(gout, dim=-1).cuda()
                         mlp_emb_ = F.log_softmax(output, dim=-1).cuda()
                         tkloss = kl_loss(mlp_emb_.cuda().float(), gout.cuda().float())
+                        
+                        # 总loss
                         loss = loss1 + 10 * tkloss + 1 * scl
-
                     else:
-                        label = target[..., : self.args["output_dim"]]
+                        # 普通训练模式
                         output, out_, _ = self.model(data)
-
-                        if self.args["real_value"]:
-                            output = self.scaler.inverse_transform(output)
-
                         loss = self.loss(output, label)
 
+                    # 反归一化
+                    d_output = self.scaler.inverse_transform(output)
+                    d_label = self.scaler.inverse_transform(label)
+
+                    # 反归一化的loss
+                    d_loss = self.loss(d_output, d_label)
+
+                    # 反向传播和优化（仅在训练模式）
                     if optimizer_step and self.optimizer is not None:
                         self.optimizer.zero_grad()
                         loss.backward()
@@ -128,20 +162,34 @@ class Trainer:
 
                     step_time = time.time() - start_time
                     self.stats.record_step_time(step_time, mode)
-                    total_loss += loss.item()
+                    total_loss += d_loss.item()
+
+                    # 累积预测结果
+                    y_pred.append(d_output.detach().cpu())
+                    y_true.append(d_label.detach().cpu())
 
                     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}"
+                            f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
                         )
 
                     # 更新 tqdm 的进度
                     pbar.update(1)
-                    pbar.set_postfix(loss=loss.item())
+                    pbar.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"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
+            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()
@@ -157,6 +205,7 @@ 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
@@ -182,13 +231,7 @@ class Trainer:
             else:
                 not_improved_count += 1
 
-            if (
+            if self._should_early_stop(not_improved_count):
-                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:
@@ -207,15 +250,26 @@ class Trainer:
         # 输出统计与参数
         self.stats.end_training()
         self.stats.report(self.logger)
-        try:
+        self._log_model_params()
-            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 _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 _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")
@@ -274,48 +328,44 @@ 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["device"])
+            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"]]
                 output, _, _ = model(data)
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
 
-        if args["real_value"]:
+        # 反归一化
-            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        else:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-            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]
+        for t in range(d_y_true.shape[1]):
-        # torch.save(y_true, "./test/PEMSD8/y_true.pt") # [3566,12,170,1]
-
-        for t in range(y_true.shape[1]):
             mae, rmse, mape = all_metrics(
-                y_pred[:, t, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

trainer/cdeTrainer/cdetrainer.py

@@ -25,37 +25,60 @@ class Trainer:
         times,
         w,
     ):
+        # 设备和基本参数
+        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 = 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
-        self.device = args["device"]
         
-        # Paths for saving models and logs
+        # 初始化路径、日志和统计
+        self._initialize_paths(train_args)
+        self._initialize_logger(train_args)
+        self._initialize_stats()
+        
+        # 模型特定参数
+        self.times = times.to(self.device, dtype=torch.float)
+        self.w = w
+    
+    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(
             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 _initialize_stats(self):
-        self.times = times.to(self.device, dtype=torch.float)
+        """初始化统计信息记录器"""
-        self.w = w
+        self.stats = TrainingStats(device=self.device)
 
     def _run_epoch(self, epoch, dataloader, mode):
+        """运行一个训练/验证/测试epoch"""
+        # 设置模型模式和是否进行优化
         if mode == "train":
             self.model.train()
             optimizer_step = True
@@ -63,24 +86,38 @@ 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):
-            with tqdm(
+            progress_bar = tqdm(
-                total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
+                enumerate(dataloader), 
-            ) as pbar:
+                total=len(dataloader), 
-                for batch_idx, batch in enumerate(dataloader):
+                desc=f"{mode.capitalize()} Epoch {epoch}"
+            )
+            
+            for batch_idx, batch in progress_bar:
                 start_time = time.time()
                 batch = tuple(b.to(self.device, dtype=torch.float) for b in batch)
                 *train_coeffs, target = batch
                 label = target[..., : self.args["output_dim"]]
+                
+                # 前向传播
                 output = self.model(self.times, train_coeffs)
                 
-                    # if self.args['real_value']:
+                # 计算原始loss
-                    #     output = self.scaler.inverse_transform(output)
-
                 loss = self.loss(output, label)
+
+                # 反归一化
+                d_output = self.scaler.inverse_transform(output)
+                d_label = self.scaler.inverse_transform(label)
+
+                # 反归一化的loss
+                d_loss = self.loss(d_output, d_label)
+
+                # 反向传播和优化（仅在训练模式）
                 if optimizer_step and self.optimizer is not None:
                     self.optimizer.zero_grad()
                     loss.backward()
@@ -91,25 +128,42 @@ class Trainer:
                         )
                     self.optimizer.step()
 
+                # 记录步骤时间
                 step_time = time.time() - start_time
                 self.stats.record_step_time(step_time, mode)
-                    total_loss += loss.item()
+                total_loss += d_loss.item()
+
+                # 累积预测结果
+                y_pred.append(d_output.detach().cpu())
+                y_true.append(d_label.detach().cpu())
 
                 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}"
+                        f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
                     )
 
                 # 更新 tqdm 的进度
-                    pbar.update(1)
+                progress_bar.update(1)
-                    pbar.set_postfix(loss=loss.item())
+                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)
-        self.logger.info(
+        
-            f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
+        # 计算并记录指标
+        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):
@@ -122,21 +176,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
@@ -145,6 +207,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._log_model_params()
+        
+        # 最终评估
+        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"]
@@ -152,30 +239,22 @@ 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}"
         )
     
-        # 输出统计与参数
+    def _log_model_params(self):
-        self.stats.end_training()
+        """输出模型可训练参数数量"""
-        self.stats.report(self.logger)
+        total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
-        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)
@@ -188,42 +267,41 @@ class Trainer:
 
     @staticmethod
     def test(model, args, data_loader, scaler, logger):
+        """对模型进行评估并输出性能指标"""
+        # 设置为评估模式
         model.eval()
+        
+        # 收集预测和真实标签
         y_pred, y_true = [], []
         times = torch.linspace(0, 11, 12)
 
+        # 不计算梯度的情况下进行预测
         with torch.no_grad():
             for batch_idx, batch in enumerate(data_loader):
-                batch = tuple(b.to(args["device"], dtype=torch.float) for b in batch)
+                batch = tuple(b.to(args["basic"]["device"], dtype=torch.float) for b in batch)
                 *test_coeffs, target = batch
                 label = target[..., : args["output_dim"]]
-                output = model(times.to(args["device"], dtype=torch.float), test_coeffs)
+                output = model(times.to(args["basic"]["device"], dtype=torch.float), test_coeffs)
-                y_true.append(label)
+                y_true.append(label.detach().cpu())
-                y_pred.append(output)
+                y_pred.append(output.detach().cpu())
 
-        # if args['real_value']:
+        # 反归一化
-        #     y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
+        d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
-        # else:
+        d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
-        y_pred = torch.cat(y_pred, dim=0)
-        y_true = torch.cat(y_true, dim=0)
 
-        for t in range(y_true.shape[1]):
+        # 计算并记录每个时间步的指标
+        for t in range(d_y_true.shape[1]):
             mae, rmse, mape = all_metrics(
-                y_pred[:, t, ...],
+                d_y_pred[:, t, ...],
-                y_true[:, t, ...],
+                d_y_true[:, t, ...],
                 args["mae_thresh"],
                 args["mape_thresh"],
             )
-            logger.info(
+            logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-                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"]
+        mae, rmse, mape = all_metrics(d_y_pred, d_y_true, args["mae_thresh"], args["mape_thresh"])
-        )
+        logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
-        logger.info(
-            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
-        )
 
     @staticmethod
     def _compute_sampling_threshold(global_step, k):

utils/training_stats.py

@@ -47,6 +47,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:
@@ -54,26 +58,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 =====")