import math import os import time import copy import torch.nn.functional as F import torch from torch import nn from tqdm import tqdm from utils.logger import get_logger from utils.loss_function import all_metrics from model.STMLP.STMLP import STMLP from utils.training_stats import TrainingStats class Trainer: def __init__( self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args, lr_scheduler=None, ): # 设备和基本参数 self.device = args["basic"]["device"] train_args = args["train"] # 模型和训练相关组件 self.model = model self.loss = loss self.optimizer = optimizer self.lr_scheduler = lr_scheduler # 数据加载器 self.train_loader = train_loader self.val_loader = val_loader self.test_loader = test_loader # 数据处理工具 self.scaler = scaler self.args = train_args # 统计信息 self.train_per_epoch = len(train_loader) self.val_per_epoch = len(val_loader) if val_loader else 0 # 初始化路径、日志和统计 self._initialize_paths(args, train_args) self._initialize_logger(train_args) self._initialize_stats() # 教师-学生蒸馏相关 if self.args["teacher_stu"]: self.tmodel = self.loadTeacher(args) else: self.logger.info( f"当前使用预训练模式，预训练后请移动教师模型到" f"./pre-train/{args['model']['type']}/{args['data']['type']}/best_model.pth" 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 else: self.model.eval() optimizer_step = False # 初始化变量 total_loss = 0 epoch_time = time.time() y_pred, y_true = [], [] with torch.set_grad_enabled(optimizer_step): with tqdm( total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" ) as pbar: for batch_idx, (data, target) in enumerate(dataloader): start_time = time.time() label = target[..., : self.args["output_dim"]] if self.args["teacher_stu"]: # 教师-学生蒸馏模式 output, out_, _ = self.model(data) gout, tout, sout = self.tmodel(data) # 计算原始loss loss1 = self.loss(output, label) # 计算蒸馏相关loss scl = self.loss_cls(out_, sout) kl_loss = nn.KLDivLoss( reduction="batchmean", log_target=True ).cuda() 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 # 检查output和label的shape是否一致 if output.shape == label.shape: print(f"✓ Test passed: output shape {output.shape} matches label shape {label.shape}") import sys sys.exit(0) else: print(f"✗ Test failed: output shape {output.shape} does not match label shape {label.shape}") import sys sys.exit(1) else: # 普通训练模式 output, out_, _ = self.model(data) loss = self.loss(output, label) # 检查output和label的shape是否一致 if output.shape == label.shape: print(f"✓ Test passed: output shape {output.shape} matches label shape {label.shape}") import sys sys.exit(0) else: print(f"✗ Test failed: output shape {output.shape} does not match label shape {label.shape}") import sys sys.exit(1) # 反归一化 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() if self.args["grad_norm"]: torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.args["max_grad_norm"] ) self.optimizer.step() step_time = time.time() - start_time self.stats.record_step_time(step_time, mode) total_loss += 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: {d_loss.item():.6f}" ) # 更新 tqdm 的进度 pbar.update(1) 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"Epoch #{epoch:02d}: {mode.capitalize():<5} MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s" ) # 记录内存 self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): return self._run_epoch(epoch, self.train_loader, "train") def val_epoch(self, epoch): return self._run_epoch(epoch, self.val_loader or self.test_loader, "val") def test_epoch(self, epoch): return self._run_epoch(epoch, self.test_loader, "test") def train(self): """执行完整的训练流程""" best_model, best_test_model = None, None best_loss, best_test_loss = float("inf"), float("inf") not_improved_count = 0 self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args["epochs"] + 1): train_epoch_loss = self.train_epoch(epoch) val_epoch_loss = self.val_epoch(epoch) test_epoch_loss = self.test_epoch(epoch) if train_epoch_loss > 1e6: self.logger.warning("Gradient explosion detected. Ending...") break if val_epoch_loss < best_loss: best_loss = val_epoch_loss not_improved_count = 0 best_model = copy.deepcopy(self.model.state_dict()) torch.save(best_model, self.best_path) torch.save(best_model, self.pretrain_path) 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()) torch.save(best_test_model, self.best_test_path) torch.save(best_model, self.pretrain_best_path) if not self.args["debug"]: torch.save(best_model, self.best_path) torch.save(best_test_model, self.best_test_path) self.logger.info( f"Best models saved at {self.best_path} and {self.best_test_path}" ) # 输出统计与参数 self.stats.end_training() self.stats.report(self.logger) 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"] ): 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") self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) self.model.load_state_dict(best_test_model) self.logger.info("Testing on best test model") self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) def loadTeacher(self, args): model_path = ( f"./pre-train/{args['model']['type']}/{args['data']['type']}/best_model.pth" ) try: # 尝试加载教师模型权重 state_dict = torch.load(model_path) self.logger.info(f"成功加载教师模型权重: {model_path}") # 初始化并返回教师模型 args["model"]["model_type"] = "teacher" tmodel = STMLP(args["model"]) tmodel = tmodel.to(args["device"]) tmodel.load_state_dict(state_dict, strict=False) return tmodel except FileNotFoundError: # 如果找不到权重文件，记录日志并修改 args self.logger.error( f"未找到教师模型权重文件: {model_path}。切换到预训练模式训练老师权重。\n" f"在预训练完成后，再次启动模型则为蒸馏模式" ) self.args["teacher_stu"] = False return None def loss_cls(self, x1, x2): temperature = 0.05 x1 = F.normalize(x1, p=2, dim=-1) x2 = F.normalize(x2, p=2, dim=-1) weight = F.cosine_similarity(x1, x2, dim=-1) batch_size = x1.size()[0] # neg score out = torch.cat([x1, x2], dim=0) neg = torch.exp( torch.matmul(out, out.transpose(2, 3).contiguous()) / temperature ) pos = torch.exp(torch.sum(x1 * x2, dim=-1) * weight / temperature) # pos = torch.exp(torch.sum(x1 * x2, dim=-1) / temperature) pos = torch.cat([pos, pos], dim=0).sum(dim=1) Ng = neg.sum(dim=-1).sum(dim=1) loss = (-torch.log(pos / (pos + Ng))).mean() return loss @staticmethod def test(model, args, data_loader, scaler, logger, path=None): """对模型进行评估并输出性能指标""" # 加载模型检查点（如果提供了路径） if path: checkpoint = torch.load(path) model.load_state_dict(checkpoint["state_dict"]) model.to(args["basic"]["device"]) # 设置为评估模式 model.eval() # 收集预测和真实标签 y_pred, y_true = [], [] # 不计算梯度的情况下进行预测 with torch.no_grad(): for data, target in data_loader: label = target[..., : args["output_dim"]] 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, ...], args["mae_thresh"], args["mape_thresh"], ) logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") # 计算并记录平均指标 mae, rmse, mape = all_metrics(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}") @staticmethod def _compute_sampling_threshold(global_step, k): return k / (k + math.exp(global_step / k))