import math import os import time import copy from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics from lib.training_stats import TrainingStats class Trainer: def __init__( self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, 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.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.logger = get_logger( train_config["log_dir"], name=self.model.__class__.__name__, debug=train_config["debug"], ) self.logger.info(f"Experiment log path in: {train_config['log_dir']}") # Stats tracker self.stats = TrainingStats(device=self.device) def _run_epoch(self, epoch, dataloader, mode): is_train = mode == "train" self.model.train() if is_train else self.model.eval() total_loss = 0.0 epoch_time = time.time() with ( torch.set_grad_enabled(is_train), tqdm( total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}" ) as pbar, ): for batch_idx, batch in enumerate(dataloader): start_time = time.time() # unpack the new cycle_index data, target, cycle_index = batch data = data.to(self.device) target = target.to(self.device) cycle_index = cycle_index.to(self.device).long() # forward if is_train: self.optimizer.zero_grad() output = self.model(data, cycle_index) else: output = self.model(data, cycle_index) # compute loss label = target[..., : self.args["output_dim"]] if self.args["real_value"]: output = self.scaler.inverse_transform(output) loss = self.loss(output, label) # backward / step if is_train: 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 += loss.item() # 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}" ) pbar.update(1) pbar.set_postfix(loss=loss.item()) 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" ) # 记录内存 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()) self.logger.info("Best validation model saved!") 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." ) 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.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) 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) @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.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) 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) # 你在这里需要把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]): mae, rmse, mape = all_metrics( y_pred[:, t, ...], 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( y_pred, 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))