Shortcuts

mmengine.logging.message_hub 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import copy
import logging
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Optional, Union

import numpy as np

from mmengine.utils import ManagerMixin
from .history_buffer import HistoryBuffer
from .logger import print_log

if TYPE_CHECKING:
    import torch


[文档]class MessageHub(ManagerMixin): """Message hub for component interaction. MessageHub is created and accessed in the same way as ManagerMixin. ``MessageHub`` will record log information and runtime information. The log information refers to the learning rate, loss, etc. of the model during training phase, which will be stored as ``HistoryBuffer``. The runtime information refers to the iter times, meta information of runner etc., which will be overwritten by next update. Args: name (str): Name of message hub used to get corresponding instance globally. log_scalars (dict, optional): Each key-value pair in the dictionary is the name of the log information such as "loss", "lr", "metric" and their corresponding values. The type of value must be HistoryBuffer. Defaults to None. runtime_info (dict, optional): Each key-value pair in the dictionary is the name of the runtime information and their corresponding values. Defaults to None. resumed_keys (dict, optional): Each key-value pair in the dictionary decides whether the key in :attr:`_log_scalars` and :attr:`_runtime_info` will be serialized. Note: Key in :attr:`_resumed_keys` belongs to :attr:`_log_scalars` or :attr:`_runtime_info`. The corresponding value cannot be set repeatedly. Examples: >>> # create empty `MessageHub`. >>> message_hub1 = MessageHub('name') >>> log_scalars = dict(loss=HistoryBuffer()) >>> runtime_info = dict(task='task') >>> resumed_keys = dict(loss=True) >>> # create `MessageHub` from data. >>> message_hub2 = MessageHub( >>> name='name', >>> log_scalars=log_scalars, >>> runtime_info=runtime_info, >>> resumed_keys=resumed_keys) """ def __init__(self, name: str, log_scalars: Optional[dict] = None, runtime_info: Optional[dict] = None, resumed_keys: Optional[dict] = None): super().__init__(name) self._log_scalars = self._parse_input('log_scalars', log_scalars) self._runtime_info = self._parse_input('runtime_info', runtime_info) self._resumed_keys = self._parse_input('resumed_keys', resumed_keys) for value in self._log_scalars.values(): assert isinstance(value, HistoryBuffer), \ ("The type of log_scalars'value must be HistoryBuffer, but " f'got {type(value)}') for key in self._resumed_keys.keys(): assert key in self._log_scalars or key in self._runtime_info, \ ('Key in `resumed_keys` must contained in `log_scalars` or ' f'`runtime_info`, but got {key}')
[文档] @classmethod def get_current_instance(cls) -> 'MessageHub': """Get latest created ``MessageHub`` instance. :obj:`MessageHub` can call :meth:`get_current_instance` before any instance has been created, and return a message hub with the instance name "mmengine". Returns: MessageHub: Empty ``MessageHub`` instance. """ if not cls._instance_dict: cls.get_instance('mmengine') return super().get_current_instance()
[文档] def update_scalar(self, key: str, value: Union[int, float, np.ndarray, 'torch.Tensor'], count: int = 1, resumed: bool = True) -> None: """Update :attr:_log_scalars. Update ``HistoryBuffer`` in :attr:`_log_scalars`. If corresponding key ``HistoryBuffer`` has been created, ``value`` and ``count`` is the argument of ``HistoryBuffer.update``, Otherwise, ``update_scalar`` will create an ``HistoryBuffer`` with value and count via the constructor of ``HistoryBuffer``. Examples: >>> message_hub = MessageHub(name='name') >>> # create loss `HistoryBuffer` with value=1, count=1 >>> message_hub.update_scalar('loss', 1) >>> # update loss `HistoryBuffer` with value >>> message_hub.update_scalar('loss', 3) >>> message_hub.update_scalar('loss', 3, resumed=False) AssertionError: loss used to be true, but got false now. resumed keys cannot be modified repeatedly' Note: The ``resumed`` argument needs to be consistent for the same ``key``. Args: key (str): Key of ``HistoryBuffer``. value (torch.Tensor or np.ndarray or int or float): Value of log. count (torch.Tensor or np.ndarray or int or float): Accumulation times of log, defaults to 1. `count` will be used in smooth statistics. resumed (str): Whether the corresponding ``HistoryBuffer`` could be resumed. Defaults to True. """ self._set_resumed_keys(key, resumed) checked_value = self._get_valid_value(value) assert isinstance(count, int), ( f'The type of count must be int. but got {type(count): {count}}') if key in self._log_scalars: self._log_scalars[key].update(checked_value, count) else: self._log_scalars[key] = HistoryBuffer([checked_value], [count])
[文档] def update_scalars(self, log_dict: dict, resumed: bool = True) -> None: """Update :attr:`_log_scalars` with a dict. ``update_scalars`` iterates through each pair of log_dict key-value, and calls ``update_scalar``. If type of value is dict, the value should be ``dict(value=xxx) or dict(value=xxx, count=xxx)``. Item in ``log_dict`` has the same resume option. Note: The ``resumed`` argument needs to be consistent for the same ``log_dict``. Args: log_dict (str): Used for batch updating :attr:`_log_scalars`. resumed (bool): Whether all ``HistoryBuffer`` referred in log_dict should be resumed. Defaults to True. Examples: >>> message_hub = MessageHub.get_instance('mmengine') >>> log_dict = dict(a=1, b=2, c=3) >>> message_hub.update_scalars(log_dict) >>> # The default count of `a`, `b` and `c` is 1. >>> log_dict = dict(a=1, b=2, c=dict(value=1, count=2)) >>> message_hub.update_scalars(log_dict) >>> # The count of `c` is 2. """ assert isinstance(log_dict, dict), ('`log_dict` must be a dict!, ' f'but got {type(log_dict)}') for log_name, log_val in log_dict.items(): if isinstance(log_val, dict): assert 'value' in log_val, \ f'value must be defined in {log_val}' count = self._get_valid_value(log_val.get('count', 1)) value = log_val['value'] else: count = 1 value = log_val assert isinstance(count, int), ('The type of count must be int. but got ' f'{type(count): {count}}') self.update_scalar(log_name, value, count, resumed)
[文档] def update_info(self, key: str, value: Any, resumed: bool = True) -> None: """Update runtime information. The key corresponding runtime information will be overwritten each time calling ``update_info``. Note: The ``resumed`` argument needs to be consistent for the same ``key``. Examples: >>> message_hub = MessageHub(name='name') >>> message_hub.update_info('iter', 100) Args: key (str): Key of runtime information. value (Any): Value of runtime information. resumed (bool): Whether the corresponding ``HistoryBuffer`` could be resumed. """ self._set_resumed_keys(key, resumed) self._runtime_info[key] = value
[文档] def pop_info(self, key: str, default: Optional[Any] = None) -> Any: """Remove runtime information by key. If the key does not exist, this method will return the default value. Args: key (str): Key of runtime information. default (Any, optional): The default returned value for the given key. Returns: Any: The runtime information if the key exists. """ return self._runtime_info.pop(key, default)
[文档] def update_info_dict(self, info_dict: dict, resumed: bool = True) -> None: """Update runtime information with dictionary. The key corresponding runtime information will be overwritten each time calling ``update_info``. Note: The ``resumed`` argument needs to be consistent for the same ``info_dict``. Examples: >>> message_hub = MessageHub(name='name') >>> message_hub.update_info({'iter': 100}) Args: info_dict (str): Runtime information dictionary. resumed (bool): Whether the corresponding ``HistoryBuffer`` could be resumed. """ assert isinstance(info_dict, dict), ('`log_dict` must be a dict!, ' f'but got {type(info_dict)}') for key, value in info_dict.items(): self.update_info(key, value, resumed=resumed)
def _set_resumed_keys(self, key: str, resumed: bool) -> None: """Set corresponding resumed keys. This method is called by ``update_scalar``, ``update_scalars`` and ``update_info`` to set the corresponding key is true or false in :attr:`_resumed_keys`. Args: key (str): Key of :attr:`_log_scalrs` or :attr:`_runtime_info`. resumed (bool): Whether the corresponding ``HistoryBuffer`` could be resumed. """ if key not in self._resumed_keys: self._resumed_keys[key] = resumed else: assert self._resumed_keys[key] == resumed, \ f'{key} used to be {self._resumed_keys[key]}, but got ' \ '{resumed} now. resumed keys cannot be modified repeatedly.' @property def log_scalars(self) -> OrderedDict: """Get all ``HistoryBuffer`` instances. Note: Considering the large memory footprint of history buffers in the post-training, :meth:`get_scalar` will return a reference of history buffer rather than a copy. Returns: OrderedDict: All ``HistoryBuffer`` instances. """ return self._log_scalars @property def runtime_info(self) -> OrderedDict: """Get all runtime information. Returns: OrderedDict: A copy of all runtime information. """ return self._runtime_info
[文档] def get_scalar(self, key: str) -> HistoryBuffer: """Get ``HistoryBuffer`` instance by key. Note: Considering the large memory footprint of history buffers in the post-training, :meth:`get_scalar` will not return a reference of history buffer rather than a copy. Args: key (str): Key of ``HistoryBuffer``. Returns: HistoryBuffer: Corresponding ``HistoryBuffer`` instance if the key exists. """ if key not in self.log_scalars: raise KeyError(f'{key} is not found in Messagehub.log_buffers: ' f'instance name is: {MessageHub.instance_name}') return self.log_scalars[key]
[文档] def get_info(self, key: str, default: Optional[Any] = None) -> Any: """Get runtime information by key. If the key does not exist, this method will return default information. Args: key (str): Key of runtime information. default (Any, optional): The default returned value for the given key. Returns: Any: A copy of corresponding runtime information if the key exists. """ if key not in self.runtime_info: return default else: # TODO: There are restrictions on objects that can be saved # return copy.deepcopy(self._runtime_info[key]) return self._runtime_info[key]
def _get_valid_value( self, value: Union['torch.Tensor', np.ndarray, np.number, int, float], ) -> Union[int, float]: """Convert value to python built-in type. Args: value (torch.Tensor or np.ndarray or np.number or int or float): value of log. Returns: float or int: python built-in type value. """ if isinstance(value, (np.ndarray, np.number)): assert value.size == 1 value = value.item() elif isinstance(value, (int, float)): value = value else: # check whether value is torch.Tensor but don't want # to import torch in this file assert hasattr(value, 'numel') and value.numel() == 1 value = value.item() return value # type: ignore
[文档] def state_dict(self) -> dict: """Returns a dictionary containing log scalars, runtime information and resumed keys, which should be resumed. The returned ``state_dict`` can be loaded by :meth:`load_state_dict`. Returns: dict: A dictionary contains ``log_scalars``, ``runtime_info`` and ``resumed_keys``. """ saved_scalars = OrderedDict() saved_info = OrderedDict() for key, value in self._log_scalars.items(): if self._resumed_keys.get(key, False): saved_scalars[key] = copy.deepcopy(value) for key, value in self._runtime_info.items(): if self._resumed_keys.get(key, False): try: saved_info[key] = copy.deepcopy(value) except: # noqa: E722 print_log( f'{key} in message_hub cannot be copied, ' f'just return its reference. ', logger='current', level=logging.WARNING) saved_info[key] = value return dict( log_scalars=saved_scalars, runtime_info=saved_info, resumed_keys=self._resumed_keys)
[文档] def load_state_dict(self, state_dict: Union['MessageHub', dict]) -> None: """Loads log scalars, runtime information and resumed keys from ``state_dict`` or ``message_hub``. If ``state_dict`` is a dictionary returned by :meth:`state_dict`, it will only make copies of data which should be resumed from the source ``message_hub``. If ``state_dict`` is a ``message_hub`` instance, it will make copies of all data from the source message_hub. We suggest to load data from ``dict`` rather than a ``MessageHub`` instance. Args: state_dict (dict or MessageHub): A dictionary contains key ``log_scalars`` ``runtime_info`` and ``resumed_keys``, or a MessageHub instance. """ if isinstance(state_dict, dict): for key in ('log_scalars', 'runtime_info', 'resumed_keys'): assert key in state_dict, ( 'The loaded `state_dict` of `MessageHub` must contain ' f'key: `{key}`') # The old `MessageHub` could save non-HistoryBuffer `log_scalars`, # therefore the loaded `log_scalars` needs to be filtered. for key, value in state_dict['log_scalars'].items(): if not isinstance(value, HistoryBuffer): print_log( f'{key} in message_hub is not HistoryBuffer, ' f'just skip resuming it.', logger='current', level=logging.WARNING) continue self.log_scalars[key] = value for key, value in state_dict['runtime_info'].items(): try: self._runtime_info[key] = copy.deepcopy(value) except: # noqa: E722 print_log( f'{key} in message_hub cannot be copied, ' f'just return its reference.', logger='current', level=logging.WARNING) self._runtime_info[key] = value for key, value in state_dict['resumed_keys'].items(): if key not in set(self.log_scalars.keys()) | \ set(self._runtime_info.keys()): print_log( f'resumed key: {key} is not defined in message_hub, ' f'just skip resuming this key.', logger='current', level=logging.WARNING) continue elif not value: print_log( f'Although resumed key: {key} is False, {key} ' 'will still be loaded this time. This key will ' 'not be saved by the next calling of ' '`MessageHub.state_dict()`', logger='current', level=logging.WARNING) self._resumed_keys[key] = value # Since some checkpoints saved serialized `message_hub` instance, # `load_state_dict` support loading `message_hub` instance for # compatibility else: self._log_scalars = copy.deepcopy(state_dict._log_scalars) self._runtime_info = copy.deepcopy(state_dict._runtime_info) self._resumed_keys = copy.deepcopy(state_dict._resumed_keys)
def _parse_input(self, name: str, value: Any) -> OrderedDict: """Parse input value. Args: name (str): name of input value. value (Any): Input value. Returns: dict: Parsed input value. """ if value is None: return OrderedDict() elif isinstance(value, dict): return OrderedDict(value) else: raise TypeError(f'{name} should be a dict or `None`, but ' f'got {type(name)}')

© Copyright 2022, mmengine contributors. Revision 39ed23fa.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.