| import os |
|
|
| import pandas as pd |
| import torch, time, wandb |
| from collections import defaultdict |
| import pytorch_lightning as pl |
| import numpy as np |
| import pdb |
| from utils.log import get_logger |
|
|
| logger = get_logger(__name__) |
|
|
|
|
|
|
|
|
|
|
| class GeneralModule(pl.LightningModule): |
| def __init__(self, args): |
| super().__init__() |
| self.save_hyperparameters() |
| self.args = args |
|
|
| self.iter_step = -1 |
| self._log = defaultdict(list) |
| self.generator = np.random.default_rng() |
| self.last_log_time = time.time() |
|
|
|
|
| def try_print_log(self): |
|
|
| step = self.iter_step if self.args.validate else self.trainer.global_step |
| if (step + 1) % self.args.print_freq == 0: |
| print(os.environ["MODEL_DIR"]) |
| log = self._log |
| log = {key: log[key] for key in log if "iter_" in key} |
|
|
| log = self.gather_log(log, self.trainer.world_size) |
| mean_log = self.get_log_mean(log) |
| mean_log.update( |
| {'epoch': float(self.trainer.current_epoch), 'step': float(self.trainer.global_step), 'iter_step': float(self.iter_step)}) |
| if self.trainer.is_global_zero: |
| print(str(mean_log)) |
| self.log_dict(mean_log, batch_size=1) |
| if self.args.wandb: |
| wandb.log(mean_log) |
| for key in list(log.keys()): |
| if "iter_" in key: |
| del self._log[key] |
|
|
| def lg(self, key, data): |
| if isinstance(data, torch.Tensor): |
| data = data.detach().cpu().item() |
| log = self._log |
| |
| if self.args.validate or self.stage == 'train': |
| log["iter_" + key].append(data) |
| log[self.stage + "_" + key].append(data) |
|
|
| def on_train_epoch_end(self): |
| log = self._log |
| log = {key: log[key] for key in log if "train_" in key} |
| log = self.gather_log(log, self.trainer.world_size) |
| mean_log = self.get_log_mean(log) |
| mean_log.update( |
| {'epoch': float(self.trainer.current_epoch), 'step': float(self.trainer.global_step), 'iter_step': float(self.iter_step)}) |
|
|
| if self.trainer.is_global_zero: |
| logger.info(str(mean_log)) |
| self.log_dict(mean_log, batch_size=1) |
| if self.args.wandb: |
| wandb.log(mean_log) |
|
|
| for key in list(log.keys()): |
| if "train_" in key: |
| del self._log[key] |
|
|
| def on_validation_epoch_end(self): |
| self.generator = np.random.default_rng() |
| log = self._log |
| log = {key: log[key] for key in log if "val_" in key} |
| log = self.gather_log(log, self.trainer.world_size) |
| mean_log = self.get_log_mean(log) |
| mean_log.update( |
| {'epoch': float(self.trainer.current_epoch), 'step': float(self.trainer.global_step), 'iter_step': float(self.iter_step)}) |
|
|
| if self.trainer.is_global_zero: |
| logger.info(str(mean_log)) |
| self.log_dict(mean_log, batch_size=1) |
| if self.args.wandb: |
| wandb.log(mean_log) |
|
|
| path = os.path.join( |
| os.environ["MODEL_DIR"], f"val_{self.trainer.global_step}.csv" |
| ) |
| pd.DataFrame(log).to_csv(path) |
|
|
| for key in list(log.keys()): |
| if "val_" in key: |
| del self._log[key] |
|
|
|
|
|
|
| def gather_log(self, log, world_size): |
| if world_size == 1: |
| return log |
| log_list = [None] * world_size |
| torch.distributed.all_gather_object(log_list, log) |
| log = {key: sum([l[key] for l in log_list], []) for key in log} |
| return log |
|
|
| def get_log_mean(self, log): |
| out = {} |
| for key in log: |
| try: |
| out[key] = np.nanmean(log[key]) |
| except: |
| pass |
| return out |
