diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..c77f418 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,4 @@ +[mypy] +explicit_package_bases = True +ignore_missing_imports = True +no_site_packages = True diff --git a/trainer/DCRNN_Trainer.py b/trainer/DCRNN_Trainer.py index 417d078..8bb2298 100755 --- a/trainer/DCRNN_Trainer.py +++ b/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.lr_scheduler = lr_scheduler + self.args = train_args + + # 统计信息 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( - total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" - ) as pbar: - for batch_idx, (data, target) in enumerate(dataloader): - start_time = time.time() - label = target[..., : self.args["output_dim"]] - output = self.model(data, labels=label.clone()).to( - self.args["device"] - ) + progress_bar = tqdm( + enumerate(dataloader), + total=len(dataloader), + desc=f"{mode.capitalize()} Epoch {epoch}" + ) + + for _, (data, target) in progress_bar: + # 记录步骤开始时间 + start_time = time.time() - if self.args["real_value"]: - output = self.scaler.inverse_transform(output) - label = self.scaler.inverse_transform(label) + # 前向传播 + label = target[..., : self.args["output_dim"]] + output = self.model(data, labels=label.clone()).to(self.device) + loss = self.loss(output, label) - loss = self.loss(output, label) - if optimizer_step and self.optimizer is not None: - self.optimizer.zero_grad() - loss.backward() + # 反归一化 + d_output = self.scaler.inverse_transform(output) + d_label = self.scaler.inverse_transform(label) - if self.args["grad_norm"]: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), self.args["max_grad_norm"] - ) - self.optimizer.step() + # 反向传播和优化（仅在训练模式） + if optimizer_step and self.optimizer is not None: + self.optimizer.zero_grad() + loss.backward() - step_time = time.time() - start_time - self.stats.record_step_time(step_time, mode) - total_loss += loss.item() - - if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0: - self.logger.info( - f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}" + # 梯度裁剪（如果需要） + 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) - # 更新 tqdm 的进度 - pbar.update(1) - pbar.set_postfix(loss=loss.item()) + # 记录步骤时间和内存使用 + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) + # 累积损失和预测结果 + total_loss += d_loss.item() + y_pred.append(d_output.detach().cpu()) + y_true.append(d_label.detach().cpu()) + + # 更新进度条 + progress_bar.set_postfix(loss=d_loss.item()) + + # 合并所有批次的预测结果 + y_pred = torch.cat(y_pred, dim=0) + y_true = torch.cat(y_true, dim=0) + + # 计算平均损失 avg_loss = total_loss / len(dataloader) - 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,38 +195,55 @@ class Trainer: else: not_improved_count += 1 - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): - self.logger.info( - f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." - ) + # 检查早停条件 + 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"]: - 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._save_best_models(best_model, best_test_model) - # 输出统计与参数 + # 结束训练并输出统计信息 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + + # 最终评估 self._finalize_training(best_model, best_test_model) + # 输出模型参数量 + self._log_model_params() + + def _should_early_stop(self, not_improved_count): + """检查是否满足早停条件""" + if ( + self.args["early_stop"] + and not_improved_count == self.args["early_stop_patience"] + ): + self.logger.info( + f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." + ) + return True + return False + + def _save_best_models(self, best_model, best_test_model): + """保存最佳模型到文件""" + torch.save(best_model, self.best_path) + torch.save(best_test_model, self.best_test_path) + self.logger.info( + f"Best models saved at {self.best_path} and {self.best_test_path}" + ) + + def _log_model_params(self): + """输出模型可训练参数数量""" + total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + + def _finalize_training(self, best_model, best_test_model): self.model.load_state_dict(best_model) self.logger.info("Testing on best validation model") @@ -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"]) - y_pred.append(output) - y_true.append(label) + output = model(data, labels=label.clone()) + y_pred.append(output.detach().cpu()) + y_true.append(label.detach().cpu()) - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) - if args["real_value"]: - y_pred = scaler.inverse_transform(y_pred) - y_true = scaler.inverse_transform(y_true) + # 反归一化 + 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(y_true.shape[1]): + # 计算并记录每个时间步的指标 + for t in range(d_y_true.shape[1]): mae, rmse, mape = all_metrics( - y_pred[:, t, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/trainer/E32Trainer.py b/trainer/E32Trainer.py index 07ff01c..5011131 100644 --- a/trainer/E32Trainer.py +++ b/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.best_test_path = os.path.join( - train_config["log_dir"], "best_test_model.pth" - ) - self.loss_figure_path = os.path.join(train_config["log_dir"], "loss.png") - - # Initialize logger - if not os.path.isdir(train_config["log_dir"]) and not train_config["debug"]: - os.makedirs(train_config["log_dir"], exist_ok=True) + # 初始化路径、日志和统计 + self._initialize_paths(train_config) + self._initialize_logger(train_config) + self._initialize_stats() + + def _initialize_paths(self, args): + """初始化模型保存路径""" + self.best_path = os.path.join(args["log_dir"], "best_model.pth") + self.best_test_path = os.path.join(args["log_dir"], "best_test_model.pth") + self.loss_figure_path = os.path.join(args["log_dir"], "loss.png") + + def _initialize_logger(self, args): + """初始化日志记录器""" + if not os.path.isdir(args["log_dir"]) and not args["debug"]: + os.makedirs(args["log_dir"], exist_ok=True) self.logger = get_logger( - 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']}") - # Stats tracker + 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""" + # 设置模型模式和是否进行优化 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,38 +200,55 @@ class Trainer: else: not_improved_count += 1 - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): - self.logger.info( - f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." - ) + # 检查早停条件 + 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"]: - 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._save_best_models(best_model, best_test_model) - # 输出统计与参数 + # 结束训练并输出统计信息 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + + # 最终评估 self._finalize_training(best_model, best_test_model) + # 输出模型参数量 + self._log_model_params() + + def _should_early_stop(self, not_improved_count): + """检查是否满足早停条件""" + if ( + self.args["early_stop"] + and not_improved_count == self.args["early_stop_patience"] + ): + self.logger.info( + f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." + ) + return True + return False + + def _save_best_models(self, best_model, best_test_model): + """保存最佳模型到文件""" + torch.save(best_model, self.best_path) + torch.save(best_test_model, self.best_test_path) + self.logger.info( + f"Best models saved at {self.best_path} and {self.best_test_path}" + ) + + def _log_model_params(self): + """输出模型可训练参数数量""" + total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + + def _finalize_training(self, best_model, best_test_model): self.model.load_state_dict(best_model) self.logger.info("Testing on best validation model") @@ -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_true.append(label) + y_pred.append(output.detach().cpu()) + y_true.append(label.detach().cpu()) - if args["real_value"]: - y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - else: - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + # 反归一化 + d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) - # 你在这里需要把y_pred和y_true保存下来 - # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1] - # torch.save(y_true, "./test/PEMS08/y_true.pt") # [3566,12,170,1] - - for t in range(y_true.shape[1]): + # 计算并记录每个时间步的指标 + for t in range(d_y_true.shape[1]): mae, rmse, mape = all_metrics( - y_pred[:, t, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/trainer/EXP_trainer.py b/trainer/EXP_trainer.py index 0416cc3..b5a48a5 100755 --- a/trainer/EXP_trainer.py +++ b/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.lr_scheduler = lr_scheduler + self.args = train_args + + # 统计信息 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( - total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" - ) as pbar: - for batch_idx, (data, target) in enumerate(dataloader): - start_time = time.time() - label = target[..., : self.args["output_dim"]] - output = self.model(data).to(self.args["device"]) + progress_bar = tqdm( + enumerate(dataloader), + total=len(dataloader), + desc=f"{mode.capitalize()} Epoch {epoch}" + ) + + for _, (data, target) in progress_bar: + # 记录步骤开始时间 + start_time = time.time() - if self.args["real_value"]: - output = self.scaler.inverse_transform(output) + # 前向传播 + label = target[..., : self.args["output_dim"]] + output = self.model(data).to(self.device) + loss = self.loss(output, label) - loss = self.loss(output, label) - if optimizer_step and self.optimizer is not None: - self.optimizer.zero_grad() - loss.backward() + # 反归一化 + d_output = self.scaler.inverse_transform(output) + d_label = self.scaler.inverse_transform(label) - if self.args["grad_norm"]: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), self.args["max_grad_norm"] - ) - self.optimizer.step() + # 反向传播和优化（仅在训练模式） + if optimizer_step and self.optimizer is not None: + self.optimizer.zero_grad() + loss.backward() - step_time = time.time() - start_time - self.stats.record_step_time(step_time, mode) - - total_loss += loss.item() - - if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0: - self.logger.info( - f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}" + # 梯度裁剪（如果需要） + 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) - # 更新 tqdm 的进度 - pbar.update(1) - pbar.set_postfix(loss=loss.item()) + # 记录步骤时间和内存使用 + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) + # 累积损失和预测结果 + total_loss += d_loss.item() + y_pred.append(d_output.detach().cpu()) + y_true.append(d_label.detach().cpu()) + + # 更新进度条 + progress_bar.set_postfix(loss=d_loss.item()) + + # 合并所有批次的预测结果 + y_pred = torch.cat(y_pred, dim=0) + y_true = torch.cat(y_true, dim=0) + + # 计算平均损失 avg_loss = total_loss / len(dataloader) - 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,37 +195,55 @@ class Trainer: else: not_improved_count += 1 - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): - self.logger.info( - f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." - ) + # 检查早停条件 + 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"]: - 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._save_best_models(best_model, best_test_model) + + # 结束训练并输出统计信息 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + + # 最终评估 self._finalize_training(best_model, best_test_model) + # 输出模型参数量 + self._log_model_params() + + def _should_early_stop(self, not_improved_count): + """检查是否满足早停条件""" + if ( + self.args["early_stop"] + and not_improved_count == self.args["early_stop_patience"] + ): + self.logger.info( + f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." + ) + return True + return False + + def _save_best_models(self, best_model, best_test_model): + """保存最佳模型到文件""" + torch.save(best_model, self.best_path) + torch.save(best_test_model, self.best_test_path) + self.logger.info( + f"Best models saved at {self.best_path} and {self.best_test_path}" + ) + + def _log_model_params(self): + """输出模型可训练参数数量""" + total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + + def _finalize_training(self, best_model, best_test_model): self.model.load_state_dict(best_model) self.logger.info("Testing on best validation model") @@ -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_true.append(label) + y_pred.append(output.detach().cpu()) + y_true.append(label.detach().cpu()) - if args["real_value"]: - y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - else: - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + # 反归一化 + d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) - # 你在这里需要把y_pred和y_true保存下来 - # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1] - # torch.save(y_true, "./test/PEMS08/y_true.pt") # [3566,12,170,1] - - for t in range(y_true.shape[1]): + # 计算并记录每个时间步的指标 + for t in range(d_y_true.shape[1]): mae, rmse, mape = all_metrics( - y_pred[:, t, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/trainer/PDG2SEQ_Trainer.py b/trainer/PDG2SEQ_Trainer.py index 72d155d..95b7a61 100755 --- a/trainer/PDG2SEQ_Trainer.py +++ b/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.lr_scheduler = lr_scheduler + self.args = train_args + self.batches_seen = 0 + + # 统计信息 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._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,55 +81,86 @@ 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( - total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" - ) as pbar: - for batch_idx, (data, target) in enumerate(dataloader): - 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"] + progress_bar = tqdm( + enumerate(dataloader), + total=len(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() + + # 前向传播 + 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() + + 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}" ) - if self.args["real_value"]: - output = self.scaler.inverse_transform(output) + # 更新 tqdm 的进度 + progress_bar.update(1) + progress_bar.set_postfix(loss=d_loss.item()) - loss = self.loss(output, label) - if optimizer_step and self.optimizer is not None: - self.optimizer.zero_grad() - loss.backward() - - if self.args["grad_norm"]: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), self.args["max_grad_norm"] - ) - self.optimizer.step() - - # record step time - step_time = time.time() - start_time - self.stats.record_step_time(step_time, mode) - total_loss += loss.item() - - if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0: - self.logger.info( - f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}" - ) - - # 更新 tqdm 的进度 - pbar.update(1) - pbar.set_postfix(loss=loss.item()) + # 合并所有批次的预测结果 + 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,37 +204,54 @@ class Trainer: else: not_improved_count += 1 - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): - self.logger.info( - f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." - ) + # 检查早停条件 + 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"]: - 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._save_best_models(best_model, best_test_model) - # 输出统计与参数 + # 结束训练并输出统计信息 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + + # 输出模型参数量 + self._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 _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) @@ -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_true.append(label) + y_pred.append(output.detach().cpu()) + y_true.append(label.detach().cpu()) - if args["real_value"]: - y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - else: - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + # 反归一化 + 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(y_true.shape[1]): + # 计算并记录每个时间步的指标 + for t in range(d_y_true.shape[1]): mae, rmse, mape = all_metrics( - y_pred[:, t, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/trainer/STMLP_Trainer.py b/trainer/STMLP_Trainer.py index 6b2217a..d1ac02a 100644 --- a/trainer/STMLP_Trainer.py +++ b/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.lr_scheduler = lr_scheduler + self.args = train_args + + # 统计信息 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.best_test_path = os.path.join(self.args["log_dir"], "best_test_model.pth") - self.loss_figure_path = os.path.join(self.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") - - # 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"]) - + # 初始化路径、日志和统计 + 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: @@ -70,9 +63,41 @@ class Trainer: 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 @@ -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() + label = target[..., : self.args["output_dim"]] + if self.args["teacher_stu"]: - label = target[..., : self.args["output_dim"]] + # 教师-学生蒸馏模式 output, out_, _ = self.model(data) gout, tout, sout = self.tmodel(data) - - if self.args["real_value"]: - output = self.scaler.inverse_transform(output) - + + # 计算原始loss 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 ( - 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." - ) + if self._should_early_stop(not_improved_count): break if test_epoch_loss < best_test_loss: @@ -207,14 +250,25 @@ class Trainer: # 输出统计与参数 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + self._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) @@ -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_true.append(label) + y_pred.append(output.detach().cpu()) + y_true.append(label.detach().cpu()) - if args["real_value"]: - y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - else: - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + # 反归一化 + d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) - # 你在这里需要把y_pred和y_true保存下来 - # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1] - # torch.save(y_true, "./test/PEMSD8/y_true.pt") # [3566,12,170,1] - - 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, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/trainer/cdeTrainer/__init__.py b/trainer/cdeTrainer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/trainer/cdeTrainer/cdetrainer.py b/trainer/cdeTrainer/cdetrainer.py index 5678a7c..84111fb 100755 --- a/trainer/cdeTrainer/cdetrainer.py +++ b/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.lr_scheduler = lr_scheduler + self.args = train_args + + # 统计信息 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"]) - self.times = times.to(self.device, dtype=torch.float) - self.w = w + + 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 @@ -63,53 +86,84 @@ 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( - total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" - ) as pbar: - for batch_idx, batch in enumerate(dataloader): - 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) + progress_bar = tqdm( + enumerate(dataloader), + total=len(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) + + # 计算原始loss + loss = self.loss(output, label) - # if self.args['real_value']: - # output = self.scaler.inverse_transform(output) + # 反归一化 + d_output = self.scaler.inverse_transform(output) + d_label = self.scaler.inverse_transform(label) - loss = self.loss(output, label) - if optimizer_step and self.optimizer is not None: - self.optimizer.zero_grad() - loss.backward() + # 反归一化的loss + d_loss = self.loss(d_output, d_label) - if self.args["grad_norm"]: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), self.args["max_grad_norm"] - ) - self.optimizer.step() + # 反向传播和优化（仅在训练模式） + if optimizer_step and self.optimizer is not None: + self.optimizer.zero_grad() + loss.backward() - step_time = time.time() - start_time - self.stats.record_step_time(step_time, mode) - total_loss += loss.item() - - if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0: - self.logger.info( - f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}" + if self.args["grad_norm"]: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.args["max_grad_norm"] ) + self.optimizer.step() - # 更新 tqdm 的进度 - pbar.update(1) - pbar.set_postfix(loss=loss.item()) + # 记录步骤时间 + 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 的进度 + progress_bar.update(1) + 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,37 +207,54 @@ class Trainer: else: not_improved_count += 1 - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): - self.logger.info( - f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." - ) + # 检查早停条件 + 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"]: - 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._save_best_models(best_model, best_test_model) - # 输出统计与参数 + # 结束训练并输出统计信息 self.stats.end_training() self.stats.report(self.logger) - try: - total_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) - self.logger.info(f"Trainable params: {total_params}") - except Exception: - pass + + # 输出模型参数量 + self._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 _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) @@ -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) - y_true.append(label) - y_pred.append(output) + output = model(times.to(args["basic"]["device"], dtype=torch.float), test_coeffs) + y_true.append(label.detach().cpu()) + y_pred.append(output.detach().cpu()) - # if args['real_value']: - # y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - # else: - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + # 反归一化 + 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(y_true.shape[1]): + # 计算并记录每个时间步的指标 + for t in range(d_y_true.shape[1]): mae, rmse, mape = all_metrics( - y_pred[:, t, ...], - y_true[:, t, ...], + 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}" - ) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - mae, rmse, mape = all_metrics( - y_pred, y_true, args["mae_thresh"], args["mape_thresh"] - ) - logger.info( - f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}" - ) + # 计算并记录平均指标 + 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): diff --git a/utils/training_stats.py b/utils/training_stats.py index ecda094..9483354 100644 --- a/utils/training_stats.py +++ b/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 = ( - sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list) - if self.gpu_mem_usage_list - else 0 - ) - 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 - ) + avg_gpu_mem = self._calculate_average(self.gpu_mem_usage_list) + avg_cpu_mem = self._calculate_average(self.cpu_mem_usage_list) + avg_train_time = self._calculate_average(self.train_time_list) + avg_infer_time = self._calculate_average(self.infer_time_list) iters_per_sec = self.total_iters / total_time if total_time > 0 else 0 logger.info("===== Training Summary =====")