logger.py

importdatetime
importlogging
importtime

from .dist_utilimportget_dist_info, master_only

initialized_logger= {}


classMessageLogger():
"""Message logger for printing.
 Args:
 opt (dict): Config. It contains the following keys:
 name (str): Exp name.
 logger (dict): Contains 'print_freq' (str) for logger interval.
 train (dict): Contains 'total_iter' (int) for total iters.
 use_tb_logger (bool): Use tensorboard logger.
 start_iter (int): Start iter. Default: 1.
 tb_logger (obj:`tb_logger`): Tensorboard logger. Default： None.
 """

def__init__(self, opt, start_iter=1, tb_logger=None):
self.exp_name=opt['name']
self.interval=opt['logger']['print_freq']
self.start_iter=start_iter
self.max_iters=opt['train']['total_iter']
self.use_tb_logger=opt['logger']['use_tb_logger']
self.tb_logger=tb_logger
self.start_time=time.time()
self.logger=get_root_logger()

@master_only
def__call__(self, log_vars):
"""Format logging message.
 Args:
 log_vars (dict): It contains the following keys:
 epoch (int): Epoch number.
 iter (int): Current iter.
 lrs (list): List for learning rates.
 time (float): Iter time.
 data_time (float): Data time for each iter.
 """
# epoch, iter, learning rates
epoch=log_vars.pop('epoch')
current_iter=log_vars.pop('iter')
lrs=log_vars.pop('lrs')

message= (f'[{self.exp_name[:5]}..][epoch:{epoch:3d}, 'f'iter:{current_iter:8,d}, lr:(')
forvinlrs:
message+=f'{v:.3e},'
message+=')] '

# time and estimated time
if'time'inlog_vars.keys():
iter_time=log_vars.pop('time')
data_time=log_vars.pop('data_time')

total_time=time.time() -self.start_time
time_sec_avg=total_time/ (current_iter-self.start_iter+1)
eta_sec=time_sec_avg* (self.max_iters-current_iter-1)
eta_str=str(datetime.timedelta(seconds=int(eta_sec)))
message+=f'[eta: {eta_str}, '
message+=f'time (data): {iter_time:.3f} ({data_time:.3f})] '

# other items, especially losses
fork, vinlog_vars.items():
message+=f'{k}: {v:.4e} '
# tensorboard logger
ifself.use_tb_logger:
# if k.startswith('l_'):
# self.tb_logger.add_scalar(f'losses/{k}', v, current_iter)
# else:
self.tb_logger.add_scalar(k, v, current_iter)
self.logger.info(message)


@master_only
definit_tb_logger(log_dir):
fromtorch.utils.tensorboardimportSummaryWriter
tb_logger=SummaryWriter(log_dir=log_dir)
returntb_logger


@master_only
definit_wandb_logger(opt):
"""We now only use wandb to sync tensorboard log."""
importwandb
logger=logging.getLogger('basicsr')

project=opt['logger']['wandb']['project']
resume_id=opt['logger']['wandb'].get('resume_id')
ifresume_id:
wandb_id=resume_id
resume='allow'
logger.warning(f'Resume wandb logger with id={wandb_id}.')
else:
wandb_id=wandb.util.generate_id()
resume='never'

wandb.init(id=wandb_id, resume=resume, name=opt['name'], config=opt, project=project, sync_tensorboard=True)

logger.info(f'Use wandb logger with id={wandb_id}; project={project}.')


defget_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
"""Get the root logger.
 The logger will be initialized if it has not been initialized. By default a
 StreamHandler will be added. If `log_file` is specified, a FileHandler will
 also be added.
 Args:
 logger_name (str): root logger name. Default: 'basicsr'.
 log_file (str | None): The log filename. If specified, a FileHandler
 will be added to the root logger.
 log_level (int): The root logger level. Note that only the process of
 rank 0 is affected, while other processes will set the level to
 "Error" and be silent most of the time.
 Returns:
 logging.Logger: The root logger.
 """
logger=logging.getLogger(logger_name)
# if the logger has been initialized, just return it
iflogger_nameininitialized_logger:
returnlogger

format_str='%(asctime)s %(levelname)s: %(message)s'
stream_handler=logging.StreamHandler()
stream_handler.setFormatter(logging.Formatter(format_str))
logger.addHandler(stream_handler)
logger.propagate=False
rank, _=get_dist_info()
ifrank!=0:
logger.setLevel('ERROR')
eliflog_fileisnotNone:
logger.setLevel(log_level)
# add file handler
# file_handler = logging.FileHandler(log_file, 'w')
file_handler=logging.FileHandler(log_file, 'a') #Shangchen: keep the previous log
file_handler.setFormatter(logging.Formatter(format_str))
file_handler.setLevel(log_level)
logger.addHandler(file_handler)
initialized_logger[logger_name] =True
returnlogger


defget_env_info():
"""Get environment information.
 Currently, only log the software version.
 """
importtorch
importtorchvision

frombasicsr.versionimport__version__
msg=r"""
 ____ _ _____ ____
 / __ ) ____ _ _____ (_)_____/ ___/ / __ \
 / __ |/ __ `// ___// // ___/\__ \ / /_/ /
 / /_/ // /_/ /(__ )/ // /__ ___/ // _, _/
 /_____/ \__,_//____//_/ \___//____//_/ |_|
 ______ __ __ __ __
 / ____/____ ____ ____/ / / / __ __ _____ / /__ / /
 / / __ / __ \ / __ \ / __ / / / / / / // ___// //_/ / /
 / /_/ // /_/ // /_/ // /_/ / / /___/ /_/ // /__ / /< /_/
 \____/ \____/ \____/ \____/ /_____/\____/ \___//_/|_| (_)
 """
msg+= ('\nVersion Information: '
f'\n\tBasicSR: {__version__}'
f'\n\tPyTorch: {torch.__version__}'
f'\n\tTorchVision: {torchvision.__version__}')
returnmsg