diff --git a/train.py b/train.py index 76ea652..acd0e60 100644 --- a/train.py +++ b/train.py @@ -6,14 +6,12 @@ import utils.initializer as init from dataloader.loader_selector import get_dataloader from trainer.trainer_selector import select_trainer -import cProfile - def read_config(config_path): with open(config_path, "r") as file: config = yaml.safe_load(file) # 全局配置 - device = "cuda:1" # 指定设备为cuda:0 + device = "cpu" # 指定设备为cuda:0 seed = 2023 # 随机种子 epochs = 120 @@ -67,8 +65,8 @@ def main(debug=False): model_list = ["iTransformer"] # 指定数据集 # dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"] - # dataset_list = ["AirQuality"] - dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"] + dataset_list = ["AirQuality"] + # dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"] # 我的调试开关，不做测试就填 str(False) # os.environ["TRY"] = str(False) @@ -99,4 +97,4 @@ def main(debug=False): if __name__ == "__main__": # 调试用 - main(debug = False) \ No newline at end of file + main(debug = True) \ No newline at end of file diff --git a/trainer/Trainer.py b/trainer/Trainer.py index 3372873..04842ba 100755 --- a/trainer/Trainer.py +++ b/trainer/Trainer.py @@ -10,47 +10,45 @@ from tqdm import tqdm class Trainer: """模型训练器，负责整个训练流程的管理""" - def __init__(self, model, loss, optimizer, - train_loader, val_loader, test_loader, scaler, - args, lr_scheduler=None,): + def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args, lr_scheduler=None): # 设备和基本参数 self.config = args self.device = args["basic"]["device"] - train_args = args["train"] + self.args = args["train"] + # 模型和训练相关组件 - self.model = model - self.loss = loss - self.optimizer = optimizer - self.lr_scheduler = lr_scheduler + self.model, self.loss, self.optimizer, self.lr_scheduler = model, loss, optimizer, lr_scheduler + # 数据加载器 - self.train_loader = train_loader - self.val_loader = val_loader - self.test_loader = test_loader + self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader + # 数据处理工具 self.scaler = scaler - self.args = train_args + # 初始化路径、日志和统计 - self._initialize_paths(train_args) - self._initialize_logger(train_args) + self._initialize_paths(self.args) + self._initialize_logger(self.args) def _initialize_paths(self, args): """初始化模型保存路径""" - self.best_path = os.path.join(args["log_dir"], "best_model.pth") - self.best_test_path = os.path.join(args["log_dir"], "best_test_model.pth") - self.loss_figure_path = os.path.join(args["log_dir"], "loss.png") + log_dir = args["log_dir"] + self.best_path = os.path.join(log_dir, "best_model.pth") + self.best_test_path = os.path.join(log_dir, "best_test_model.pth") + self.loss_figure_path = os.path.join(log_dir, "loss.png") def _initialize_logger(self, args): """初始化日志记录器""" - if not os.path.isdir(args["log_dir"]) and not args["debug"]: - os.makedirs(args["log_dir"], exist_ok=True) - self.logger = get_logger(args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]) - self.logger.info(f"Experiment log path in: {args['log_dir']}") + log_dir = args["log_dir"] + if not args["debug"]: + os.makedirs(log_dir, exist_ok=True) + self.logger = get_logger(log_dir, name=self.model.__class__.__name__, debug=args["debug"]) + self.logger.info(f"Experiment log path in: {log_dir}") def _run_epoch(self, epoch, dataloader, mode): """运行一个训练/验证/测试epoch""" # 设置模型模式和是否进行优化 - if mode == "train": self.model.train(); optimizer_step = True - else: self.model.eval(); optimizer_step = False + self.model.train() if mode == "train" else self.model.eval() + optimizer_step = mode == "train" # 初始化变量 total_loss = 0 @@ -60,105 +58,111 @@ class Trainer: # 训练/验证循环 with torch.set_grad_enabled(optimizer_step): progress_bar = tqdm( - enumerate(dataloader), + dataloader, total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" ) - for _, (data, target) in progress_bar: - # 转移数据 - data = data.to(self.device) - target = target.to(self.device) + for data, target in progress_bar: + # 转移数据并提取标签 + data, target = data.to(self.device), target.to(self.device) label = target[..., : self.args["output_dim"]] - # 计算loss和反归一化loss + + # 计算输出 output = self.model(data) + # 我的调试开关 if os.environ.get("TRY") == "True": - print(f"[{'✅' if output.shape == label.shape else '❌'}]: output: {output.shape}, label: {label.shape}") + status = '✅' if output.shape == label.shape else '❌' + print(f"[{status}]: output: {output.shape}, label: {label.shape}") assert False + + # 计算损失 loss = self.loss(output, label) d_output = self.scaler.inverse_transform(output) d_label = self.scaler.inverse_transform(label) d_loss = self.loss(d_output, d_label) + # 累积损失和预测结果 total_loss += d_loss.item() y_pred.append(d_output.detach().cpu()) y_true.append(d_label.detach().cpu()) + # 反向传播和优化（仅在训练模式） if optimizer_step and self.optimizer is not None: self.optimizer.zero_grad() loss.backward() + # 梯度裁剪（如果需要） if self.args["grad_norm"]: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args["max_grad_norm"]) + self.optimizer.step() + # 更新进度条 progress_bar.set_postfix(loss=d_loss.item()) # 合并所有批次的预测结果 - y_pred = torch.cat(y_pred, dim=0) - y_true = torch.cat(y_true, dim=0) + y_pred, y_true = torch.cat(y_pred, dim=0), torch.cat(y_true, dim=0) + # 计算损失并记录指标 avg_loss = total_loss / len(dataloader) mae, rmse, mape = all_metrics(y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]) + self.logger.info( f"Epoch #{epoch:02d}: {mode.capitalize():<5} " f"MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s" ) + return avg_loss - def train_epoch(self, epoch): - return self._run_epoch(epoch, self.train_loader, "train") - - def val_epoch(self, epoch): - return self._run_epoch(epoch, self.val_loader or self.test_loader, "val") - - def test_epoch(self, epoch): - return self._run_epoch(epoch, self.test_loader, "test") - def train(self): # 初始化记录 - best_model, best_test_model = None, None - best_loss, best_test_loss = float("inf"), float("inf") + best_model = best_test_model = None + best_loss = best_test_loss = float("inf") not_improved_count = 0 + # 开始训练 self.logger.info("Training process started") + # 训练循环 for epoch in range(1, self.args["epochs"] + 1): # 训练、验证和测试一个epoch - train_epoch_loss = self.train_epoch(epoch) - val_epoch_loss = self.val_epoch(epoch) - test_epoch_loss = self.test_epoch(epoch) + train_epoch_loss = self._run_epoch(epoch, self.train_loader, "train") + val_epoch_loss = self._run_epoch(epoch, self.val_loader or self.test_loader, "val") + test_epoch_loss = self._run_epoch(epoch, self.test_loader, "test") + # 检查梯度爆炸 if train_epoch_loss > 1e6: self.logger.warning("Gradient explosion detected. Ending...") break + # 更新最佳验证模型 if val_epoch_loss < best_loss: - best_loss = val_epoch_loss - not_improved_count = 0 + best_loss, not_improved_count = val_epoch_loss, 0 best_model = copy.deepcopy(self.model.state_dict()) self.logger.info("Best validation model saved!") else: not_improved_count += 1 - # 早停 + + # 早停检查 if self._should_early_stop(not_improved_count): break + # 更新最佳测试模型 if test_epoch_loss < best_test_loss: best_test_loss = test_epoch_loss best_test_model = copy.deepcopy(self.model.state_dict()) + # 保存最佳模型 if not self.args["debug"]: self._save_best_models(best_model, best_test_model) + # 最终评估 self._finalize_training(best_model, best_test_model) def _should_early_stop(self, not_improved_count): """检查是否满足早停条件""" - if ( - self.args["early_stop"] - and not_improved_count == self.args["early_stop_patience"] - ): + if self.args["early_stop"] and not_improved_count == self.args["early_stop_patience"]: self.logger.info( f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." ) @@ -190,58 +194,43 @@ class Trainer: @staticmethod def test(model, args, data_loader, scaler, logger, path=None): """对模型进行评估并输出性能指标""" - # 确定设备信息 - device = None - output_dim = None - # 处理不同的参数格式 - if isinstance(args, dict): - if "basic" in args: - # 完整配置情况 - device = args["basic"]["device"] - output_dim = args["train"]["output_dim"] - else: - # 只有train_args情况，从模型获取设备 - device = next(model.parameters()).device - output_dim = args["output_dim"] - else: + # 验证参数类型 + if not isinstance(args, dict): raise ValueError(f"Unsupported args type: {type(args)}") + # 确定设备和输出维度 + is_full_config = "basic" in args + device = args["basic"]["device"] if is_full_config else next(model.parameters()).device + output_dim = args["train"]["output_dim"] if is_full_config else args["output_dim"] + + # 获取metrics参数 + train_args = args["train"] if is_full_config else args + mae_thresh, mape_thresh = train_args["mae_thresh"], train_args["mape_thresh"] + # 加载模型检查点（如果提供了路径） if path: checkpoint = torch.load(path) model.load_state_dict(checkpoint["state_dict"]) model.to(device) - # 设置为评估模式 + # 设置为评估模式并收集预测结果 model.eval() - - # 收集预测和真实标签 y_pred, y_true = [], [] # 不计算梯度的情况下进行预测 with torch.no_grad(): for data, target in data_loader: # 将数据和标签移动到指定设备 - data = data.to(device) - target = target.to(device) - + data, target = data.to(device), target.to(device) label = target[..., : output_dim] + output = model(data) y_pred.append(output.detach().cpu()) y_true.append(label.detach().cpu()) + # 反归一化并计算指标 d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) - - # 获取metrics参数 - if "basic" in args: - # 完整配置情况 - mae_thresh = args["train"]["mae_thresh"] - mape_thresh = args["train"]["mape_thresh"] - else: - # 只有train_args情况 - mae_thresh = args["mae_thresh"] - mape_thresh = args["mape_thresh"] # 计算并记录每个时间步的指标 for t in range(d_y_true.shape[1]): @@ -254,9 +243,5 @@ class Trainer: logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") # 计算并记录平均指标 - mae, rmse, mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh) - logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") - - @staticmethod - def _compute_sampling_threshold(global_step, k): - return k / (k + math.exp(global_step / k)) + avg_mae, avg_rmse, avg_mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh) + logger.info(f"Average Horizon, MAE: {avg_mae:.4f}, RMSE: {avg_rmse:.4f}, MAPE: {avg_mape:.4f}")