|
|
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 |
