TrafficWheel/trainer/Trainer.py

import math
import os
import time
import copy
import torch
from utils.logger import get_logger
from utils.loss_function import all_metrics
from tqdm import tqdm

class Trainer:
    """模型训练器，负责整个训练流程的管理"""
    
    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"]
        self.args = args["train"]
        
        # 模型和训练相关组件
        self.model, self.loss, self.optimizer, self.lr_scheduler = model, loss, optimizer, lr_scheduler

        # 数据加载器
        self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader
        
        # 数据处理工具
        self.scaler = scaler
        
        # 初始化路径、日志和统计
        self._initialize_paths(self.args)
        self._initialize_logger(self.args)
    
    def _initialize_paths(self, args):
        """初始化模型保存路径"""
        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):
        """初始化日志记录器"""
        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"""
        # 设置模型模式和是否进行优化
        self.model.train() if mode == "train" else self.model.eval()
        optimizer_step = mode == "train"

        # 初始化变量
        total_loss = 0
        epoch_time = time.time()
        y_pred, y_true = [], []

        # 训练/验证循环
        with torch.set_grad_enabled(optimizer_step):
            progress_bar = tqdm(
                dataloader, 
                total=len(dataloader), 
                desc=f"{mode.capitalize()} Epoch {epoch}"
            )
            for data, target in progress_bar:
                # 转移数据并提取标签
                data, target = data.to(self.device), target.to(self.device)
                label = target[..., : self.args["output_dim"]]
                
                # 计算输出
                output = self.model(data)
                
                # 我的调试开关
                if os.environ.get("TRY") == "True":
                    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, 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(self):
        # 初始化记录
        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._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, 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"]:
            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):
        """对模型进行评估并输出性能指标"""
        # 验证参数类型
        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, 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))
        
        # 计算并记录每个时间步的指标
        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}")

        # 计算并记录平均指标
        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}")