|
- import logging
- import torch.distributed as dist
-
- import copy
- import logging
- import os
- from collections import defaultdict
- import torch
- import torch.nn as nn
-
- from typing import Any
- from typing import Optional, List, Dict, NamedTuple, Tuple, Iterable
-
- from termcolor import colored
-
- logger_initialized = {}
-
- def get_root_logger(log_file=None, log_level=logging.INFO, name='main'):
- """Get root logger and add a keyword filter to it.
- 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. The name of the root logger is the top-level package name,
- e.g., "mmdet3d".
- Args:
- log_file (str, optional): File path of log. Defaults to None.
- log_level (int, optional): The level of logger.
- Defaults to logging.INFO.
- name (str, optional): The name of the root logger, also used as a
- filter keyword. Defaults to 'mmdet3d'.
- Returns:
- :obj:`logging.Logger`: The obtained logger
- """
- logger = get_logger(name=name, log_file=log_file, log_level=log_level)
- # add a logging filter
- logging_filter = logging.Filter(name)
- logging_filter.filter = lambda record: record.find(name) != -1
-
- return logger
-
-
- 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
- # handle hierarchical names
- # e.g., logger "a" is initialized, then logger "a.b" will skip the
- # initialization since it is a child of "a".
- for logger_name in logger_initialized:
- if name.startswith(logger_name):
- return logger
-
- # handle duplicate logs to the console
- # Starting in 1.8.0, PyTorch DDP attaches a StreamHandler <stderr> (NOTSET)
- # to the root logger. As logger.propagate is True by default, this root
- # level handler causes logging messages from rank>0 processes to
- # unexpectedly show up on the console, creating much unwanted clutter.
- # To fix this issue, we set the root logger's StreamHandler, if any, to log
- # at the ERROR level.
- 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
-
- # only rank 0 will add a FileHandler
- if rank == 0 and log_file is not None:
- # Here, the default behaviour of the official logger is 'a'. Thus, we
- # provide an interface to change the file mode to the default
- # behaviour.
- 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)}')
-
- def get_missing_parameters_message(keys: List[str]) -> str:
- """
- Get a logging-friendly message to report parameter names (keys) that are in
- the model but not found in a checkpoint.
- Args:
- keys (list[str]): List of keys that were not found in the checkpoint.
- Returns:
- str: message.
- """
- groups = _group_checkpoint_keys(keys)
- msg = "Some model parameters or buffers are not found in the checkpoint:\n"
- msg += "\n".join(
- " " + colored(k + _group_to_str(v), "blue") for k, v in groups.items()
- )
- return msg
-
-
- def get_unexpected_parameters_message(keys: List[str]) -> str:
- """
- Get a logging-friendly message to report parameter names (keys) that are in
- the checkpoint but not found in the model.
- Args:
- keys (list[str]): List of keys that were not found in the model.
- Returns:
- str: message.
- """
- groups = _group_checkpoint_keys(keys)
- msg = "The checkpoint state_dict contains keys that are not used by the model:\n"
- msg += "\n".join(
- " " + colored(k + _group_to_str(v), "magenta") for k, v in groups.items()
- )
- return msg
-
-
- def _strip_prefix_if_present(state_dict: Dict[str, Any], prefix: str) -> None:
- """
- Strip the prefix in metadata, if any.
- Args:
- state_dict (OrderedDict): a state-dict to be loaded to the model.
- prefix (str): prefix.
- """
- keys = sorted(state_dict.keys())
- if not all(len(key) == 0 or key.startswith(prefix) for key in keys):
- return
-
- for key in keys:
- newkey = key[len(prefix):]
- state_dict[newkey] = state_dict.pop(key)
-
- # also strip the prefix in metadata, if any..
- try:
- metadata = state_dict._metadata # pyre-ignore
- except AttributeError:
- pass
- else:
- for key in list(metadata.keys()):
- # for the metadata dict, the key can be:
- # '': for the DDP module, which we want to remove.
- # 'module': for the actual model.
- # 'module.xx.xx': for the rest.
-
- if len(key) == 0:
- continue
- newkey = key[len(prefix):]
- metadata[newkey] = metadata.pop(key)
-
-
- def _group_checkpoint_keys(keys: List[str]) -> Dict[str, List[str]]:
- """
- Group keys based on common prefixes. A prefix is the string up to the final
- "." in each key.
- Args:
- keys (list[str]): list of parameter names, i.e. keys in the model
- checkpoint dict.
- Returns:
- dict[list]: keys with common prefixes are grouped into lists.
- """
- groups = defaultdict(list)
- for key in keys:
- pos = key.rfind(".")
- if pos >= 0:
- head, tail = key[:pos], [key[pos + 1:]]
- else:
- head, tail = key, []
- groups[head].extend(tail)
- return groups
-
-
- def _group_to_str(group: List[str]) -> str:
- """
- Format a group of parameter name suffixes into a loggable string.
- Args:
- group (list[str]): list of parameter name suffixes.
- Returns:
- str: formated string.
- """
- if len(group) == 0:
- return ""
-
- if len(group) == 1:
- return "." + group[0]
-
- return ".{" + ", ".join(group) + "}"
-
-
- def _named_modules_with_dup(
- model: nn.Module, prefix: str = ""
- ) -> Iterable[Tuple[str, nn.Module]]:
- """
- The same as `model.named_modules()`, except that it includes
- duplicated modules that have more than one name.
- """
- yield prefix, model
- for name, module in model._modules.items(): # pyre-ignore
- if module is None:
- continue
- submodule_prefix = prefix + ("." if prefix else "") + name
- yield from _named_modules_with_dup(module, submodule_prefix)
|