|
|
|
import logging |
|
|
|
import torch.distributed as dist |
|
|
|
logger_initialized = {} |
|
|
|
|
|
def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'): |
|
"""Initialize and get a logger by name. |
|
|
|
If the logger has not been initialized, this method will initialize the |
|
logger by adding one or two handlers, otherwise the initialized logger will |
|
be directly returned. During initialization, a StreamHandler will always be |
|
added. If `log_file` is specified and the process rank is 0, a FileHandler |
|
will also be added. |
|
|
|
Args: |
|
name (str): Logger name. |
|
log_file (str | None): The log filename. If specified, a FileHandler |
|
will be added to the logger. |
|
log_level (int): The logger level. Note that only the process of |
|
rank 0 is affected, and other processes will set the level to |
|
"Error" thus be silent most of the time. |
|
file_mode (str): The file mode used in opening log file. |
|
Defaults to 'w'. |
|
|
|
Returns: |
|
logging.Logger: The expected logger. |
|
""" |
|
logger = logging.getLogger(name) |
|
if name in logger_initialized: |
|
return logger |
|
|
|
|
|
|
|
for logger_name in logger_initialized: |
|
if name.startswith(logger_name): |
|
return logger |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for handler in logger.root.handlers: |
|
if type(handler) is logging.StreamHandler: |
|
handler.setLevel(logging.ERROR) |
|
|
|
stream_handler = logging.StreamHandler() |
|
handlers = [stream_handler] |
|
|
|
if dist.is_available() and dist.is_initialized(): |
|
rank = dist.get_rank() |
|
else: |
|
rank = 0 |
|
|
|
|
|
if rank == 0 and log_file is not None: |
|
|
|
|
|
|
|
file_handler = logging.FileHandler(log_file, file_mode) |
|
handlers.append(file_handler) |
|
|
|
formatter = logging.Formatter( |
|
'%(asctime)s - %(name)s - %(levelname)s - %(message)s') |
|
for handler in handlers: |
|
handler.setFormatter(formatter) |
|
handler.setLevel(log_level) |
|
logger.addHandler(handler) |
|
|
|
if rank == 0: |
|
logger.setLevel(log_level) |
|
else: |
|
logger.setLevel(logging.ERROR) |
|
|
|
logger_initialized[name] = True |
|
|
|
return logger |
|
|
|
|
|
def print_log(msg, logger=None, level=logging.INFO): |
|
"""Print a log message. |
|
|
|
Args: |
|
msg (str): The message to be logged. |
|
logger (logging.Logger | str | None): The logger to be used. |
|
Some special loggers are: |
|
- "silent": no message will be printed. |
|
- other str: the logger obtained with `get_root_logger(logger)`. |
|
- None: The `print()` method will be used to print log messages. |
|
level (int): Logging level. Only available when `logger` is a Logger |
|
object or "root". |
|
""" |
|
if logger is None: |
|
print(msg) |
|
elif isinstance(logger, logging.Logger): |
|
logger.log(level, msg) |
|
elif logger == 'silent': |
|
pass |
|
elif isinstance(logger, str): |
|
_logger = get_logger(logger) |
|
_logger.log(level, msg) |
|
else: |
|
raise TypeError( |
|
'logger should be either a logging.Logger object, str, ' |
|
f'"silent" or None, but got {type(logger)}') |
|
|