from pathlib import Path import os import time import logging from omegaconf import OmegaConf from pytorch_lightning.utilities.rank_zero import rank_zero_only def create_logger(cfg, phase='train'): # root dir set by cfg root_output_dir = Path(cfg.FOLDER) # set up logger if not root_output_dir.exists(): print('=> creating {}'.format(root_output_dir)) root_output_dir.mkdir() cfg_name = cfg.NAME model = cfg.model.target.split('.')[-2] cfg_name = os.path.basename(cfg_name).split('.')[0] final_output_dir = root_output_dir / model / cfg_name cfg.FOLDER_EXP = str(final_output_dir) time_str = time.strftime('%Y-%m-%d-%H-%M-%S') new_dir(cfg, phase, time_str, final_output_dir) head = '%(asctime)-15s %(message)s' logger = config_logger(final_output_dir, time_str, phase, head) if logger is None: logger = logging.getLogger() logger.setLevel(logging.CRITICAL) logging.basicConfig(format=head) return logger @rank_zero_only def config_logger(final_output_dir, time_str, phase, head): log_file = '{}_{}_{}.log'.format('log', time_str, phase) final_log_file = final_output_dir / log_file logging.basicConfig(filename=str(final_log_file)) logger = logging.getLogger() logger.setLevel(logging.INFO) console = logging.StreamHandler() formatter = logging.Formatter(head) console.setFormatter(formatter) logging.getLogger('').addHandler(console) file_handler = logging.FileHandler(final_log_file, 'w') file_handler.setFormatter(logging.Formatter(head)) file_handler.setLevel(logging.INFO) logging.getLogger('').addHandler(file_handler) return logger @rank_zero_only def new_dir(cfg, phase, time_str, final_output_dir): # new experiment folder cfg.TIME = str(time_str) if os.path.exists(final_output_dir) and not os.path.exists(cfg.TRAIN.RESUME) and not cfg.DEBUG and phase not in ['test', 'demo']: file_list = sorted(os.listdir(final_output_dir), reverse=True) for item in file_list: if item.endswith('.log'): os.rename(str(final_output_dir), str(final_output_dir) + '_' + cfg.TIME) break final_output_dir.mkdir(parents=True, exist_ok=True) # write config yaml config_file = '{}_{}_{}.yaml'.format('config', time_str, phase) final_config_file = final_output_dir / config_file OmegaConf.save(config=cfg, f=final_config_file)