import torch import math import os import time import copy import numpy as np # import pynvml from lib.logger import get_logger from lib.loss_function import all_metrics class Trainer(object): def __init__( self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args, lr_scheduler=None, ): super(Trainer, self).__init__() 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 = args self.lr_scheduler = lr_scheduler self.train_per_epoch = len(train_loader) if val_loader != None: self.val_per_epoch = len(val_loader) 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") # log if os.path.isdir(args["log_dir"]) == False 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("Experiment log path in: {}".format(args["log_dir"])) def val_epoch(self, epoch, val_dataloader): self.model.eval() total_val_loss = 0 epoch_time = time.time() with torch.no_grad(): for batch_idx, (data, target) in enumerate(val_dataloader): data = data label = target[..., : self.args["output_dim"]] output = self.model(data) if self.args["real_value"]: output = self.scaler.inverse_transform(output) loss = self.loss(output.cuda(), label) if not torch.isnan(loss): total_val_loss += loss.item() val_loss = total_val_loss / len(val_dataloader) self.logger.info( "Val Epoch {}: average Loss: {:.6f}, train time: {:.2f} s".format( epoch, val_loss, time.time() - epoch_time ) ) return val_loss def test_epoch(self, epoch, test_dataloader): self.model.eval() total_test_loss = 0 epoch_time = time.time() with torch.no_grad(): for batch_idx, (data, target) in enumerate(test_dataloader): data = data label = target[..., : self.args["output_dim"]] output = self.model(data) if self.args["real_value"]: output = self.scaler.inverse_transform(output) loss = self.loss(output.cuda(), label) if not torch.isnan(loss): total_test_loss += loss.item() test_loss = total_test_loss / len(test_dataloader) self.logger.info( "test Epoch {}: average Loss: {:.6f}, train time: {:.2f} s".format( epoch, test_loss, time.time() - epoch_time ) ) return test_loss def train_epoch(self, epoch): self.model.train() total_loss = 0 epoch_time = time.time() for batch_idx, (data, target) in enumerate(self.train_loader): data = data label = target[..., : self.args["output_dim"]] self.optimizer.zero_grad() output = self.model(data) if self.args["real_value"]: output = self.scaler.inverse_transform(output) loss = self.loss(output.cuda(), label) loss.backward() # add max grad clipping if self.args["grad_norm"]: torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.args["max_grad_norm"] ) self.optimizer.step() total_loss += loss.item() # log information if (batch_idx + 1) % self.args["log_step"] == 0: self.logger.info( "Train Epoch {}: {}/{} Loss: {:.6f}".format( epoch, batch_idx + 1, self.train_per_epoch, loss.item() ) ) train_epoch_loss = total_loss / self.train_per_epoch self.logger.info( "Train Epoch {}: averaged Loss: {:.6f}, train time: {:.2f} s".format( epoch, train_epoch_loss, time.time() - epoch_time ) ) # learning rate decay if self.args["lr_decay"]: self.lr_scheduler.step() return train_epoch_loss def train(self): best_model = None best_test_model = None not_improved_count = 0 best_loss = float("inf") best_test_loss = float("inf") vaild_loss = [] for epoch in range(0, self.args["epochs"]): train_epoch_loss = self.train_epoch(epoch) if self.val_loader == None: val_dataloader = self.test_loader else: val_dataloader = self.val_loader test_dataloader = self.test_loader val_epoch_loss = self.val_epoch(epoch, val_dataloader) vaild_loss.append(val_epoch_loss) test_epoch_loss = self.test_epoch(epoch, test_dataloader) 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_state = True else: not_improved_count += 1 best_state = False # early stop if self.args["early_stop"]: if not_improved_count == self.args["early_stop_patience"]: self.logger.info( "Validation performance didn't improve for {} epochs. " "Training stops.".format(self.args["early_stop_patience"]) ) break # save the best state if best_state == True: self.logger.info("Current best model saved!") best_model = copy.deepcopy(self.model.state_dict()) if test_epoch_loss < best_test_loss: best_test_loss = test_epoch_loss best_test_model = copy.deepcopy(self.model.state_dict()) # save the best model to file if not self.args["debug"]: torch.save(best_model, self.best_path) self.logger.info("Saving current best model to " + self.best_path) torch.save(best_test_model, self.best_test_path) self.logger.info("Saving current best model to " + self.best_test_path) # test self.model.load_state_dict(best_model) self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) self.logger.info("This is best_test_model") self.model.load_state_dict(best_test_model) self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) def save_checkpoint(self): state = { "state_dict": self.model.state_dict(), "optimizer": self.optimizer.state_dict(), "config": self.args, } torch.save(state, self.best_path) self.logger.info("Saving current best model to " + self.best_path) @staticmethod def test(model, args, data_loader, scaler, logger, path=None): if path != None: check_point = torch.load(path) state_dict = check_point["state_dict"] args = check_point["config"] model.load_state_dict(state_dict) model.to(args["device"]) model.eval() y_pred = [] y_true = [] with torch.no_grad(): for batch_idx, (data, target) in enumerate(data_loader): data = data label = target[..., : args["output_dim"]] output = model(data) y_true.append(label) y_pred.append(output) if args["real_value"]: y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) y_true = torch.cat(y_true, dim=0) else: y_pred = torch.cat(y_pred, dim=0) y_true = torch.cat(y_true, dim=0) 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( "Horizon {:02d}, MAE: {:.4f}, RMSE: {:.4f}, MAPE: {:.4f}".format( t + 1, mae, rmse, mape ) ) mae, rmse, mape = all_metrics( y_pred, y_true, args["mae_thresh"], args["mape_thresh"] ) logger.info( "Average Horizon, MAE: {:.4f}, RMSE: {:.4f}, MAPE: {:.4f}".format( mae, rmse, mape ) ) @staticmethod def _compute_sampling_threshold(global_step, k): """ Computes the sampling probability for scheduled sampling using inverse sigmoid. :param global_step: :param k: :return: """ return k / (k + math.exp(global_step / k))