| import time | |
| import wandb | |
| import shutil | |
| import logging | |
| import os.path as osp | |
| from torch.utils.tensorboard import SummaryWriter | |
| def mv_archived_logger(name): | |
| timestamp = time.strftime("%Y-%m-%d_%H:%M:%S_", time.localtime()) | |
| basename = 'archived_' + timestamp + osp.basename(name) | |
| archived_name = osp.join(osp.dirname(name), basename) | |
| shutil.move(name, archived_name) | |
| class CustomLogger: | |
| def __init__(self, common_cfg, tb_cfg=None, wandb_cfg=None, rank=0): | |
| global global_logger | |
| self.rank = rank | |
| if self.rank == 0: | |
| self.logger = logging.getLogger('VFI') | |
| self.logger.setLevel(logging.INFO) | |
| format_str = logging.Formatter(common_cfg['format']) | |
| console_handler = logging.StreamHandler() | |
| console_handler.setFormatter(format_str) | |
| if osp.exists(common_cfg['filename']): | |
| mv_archived_logger(common_cfg['filename']) | |
| file_handler = logging.FileHandler(common_cfg['filename'], | |
| common_cfg['filemode']) | |
| file_handler.setFormatter(format_str) | |
| self.logger.addHandler(console_handler) | |
| self.logger.addHandler(file_handler) | |
| self.tb_logger = None | |
| self.enable_wandb = False | |
| if wandb_cfg is not None: | |
| self.enable_wandb = True | |
| wandb.init(**wandb_cfg) | |
| if tb_cfg is not None: | |
| self.tb_logger = SummaryWriter(**tb_cfg) | |
| global_logger = self | |
| def __call__(self, msg=None, level=logging.INFO, tb_msg=None): | |
| if self.rank != 0: | |
| return | |
| if msg is not None: | |
| self.logger.log(level, msg) | |
| if self.tb_logger is not None and tb_msg is not None: | |
| self.tb_logger.add_scalar(*tb_msg) | |
| def close(self): | |
| if self.rank == 0 and self.enable_wandb: | |
| wandb.finish() | |