File size: 3,846 Bytes
3527383 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
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
# pdb.set_trace()
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 |