| | |
| | import warnings |
| | from typing import Any, Callable, Optional, Sequence, Tuple, Union |
| |
|
| | import numpy as np |
| |
|
| |
|
| | class HistoryBuffer: |
| | """Unified storage format for different log types. |
| | |
| | ``HistoryBuffer`` records the history of log for further statistics. |
| | |
| | Examples: |
| | >>> history_buffer = HistoryBuffer() |
| | >>> # Update history_buffer. |
| | >>> history_buffer.update(1) |
| | >>> history_buffer.update(2) |
| | >>> history_buffer.min() # minimum of (1, 2) |
| | 1 |
| | >>> history_buffer.max() # maximum of (1, 2) |
| | 2 |
| | >>> history_buffer.mean() # mean of (1, 2) |
| | 1.5 |
| | >>> history_buffer.statistics('mean') # access method by string. |
| | 1.5 |
| | |
| | Args: |
| | log_history (Sequence): History logs. Defaults to []. |
| | count_history (Sequence): Counts of history logs. Defaults to []. |
| | max_length (int): The max length of history logs. Defaults to 1000000. |
| | """ |
| | _statistics_methods: dict = dict() |
| |
|
| | def __init__(self, |
| | log_history: Sequence = [], |
| | count_history: Sequence = [], |
| | max_length: int = 1000000): |
| |
|
| | self.max_length = max_length |
| | self._set_default_statistics() |
| | assert len(log_history) == len(count_history), \ |
| | 'The lengths of log_history and count_histroy should be equal' |
| | if len(log_history) > max_length: |
| | warnings.warn(f'The length of history buffer({len(log_history)}) ' |
| | f'exceeds the max_length({max_length}), the first ' |
| | 'few elements will be ignored.') |
| | self._log_history = np.array(log_history[-max_length:]) |
| | self._count_history = np.array(count_history[-max_length:]) |
| | else: |
| | self._log_history = np.array(log_history) |
| | self._count_history = np.array(count_history) |
| |
|
| | def _set_default_statistics(self) -> None: |
| | """Register default statistic methods: min, max, current and mean.""" |
| | self._statistics_methods.setdefault('min', HistoryBuffer.min) |
| | self._statistics_methods.setdefault('max', HistoryBuffer.max) |
| | self._statistics_methods.setdefault('current', HistoryBuffer.current) |
| | self._statistics_methods.setdefault('mean', HistoryBuffer.mean) |
| |
|
| | def update(self, log_val: Union[int, float], count: int = 1) -> None: |
| | """update the log history. |
| | |
| | If the length of the buffer exceeds ``self._max_length``, the oldest |
| | element will be removed from the buffer. |
| | |
| | Args: |
| | log_val (int or float): The value of log. |
| | count (int): The accumulation times of log, defaults to 1. |
| | ``count`` will be used in smooth statistics. |
| | """ |
| | if (not isinstance(log_val, (int, float)) |
| | or not isinstance(count, (int, float))): |
| | raise TypeError(f'log_val must be int or float but got ' |
| | f'{type(log_val)}, count must be int but got ' |
| | f'{type(count)}') |
| | self._log_history = np.append(self._log_history, log_val) |
| | self._count_history = np.append(self._count_history, count) |
| | if len(self._log_history) > self.max_length: |
| | self._log_history = self._log_history[-self.max_length:] |
| | self._count_history = self._count_history[-self.max_length:] |
| |
|
| | @property |
| | def data(self) -> Tuple[np.ndarray, np.ndarray]: |
| | """Get the ``_log_history`` and ``_count_history``. |
| | |
| | Returns: |
| | Tuple[np.ndarray, np.ndarray]: History logs and the counts of |
| | the history logs. |
| | """ |
| | return self._log_history, self._count_history |
| |
|
| | @classmethod |
| | def register_statistics(cls, method: Callable) -> Callable: |
| | """Register custom statistics method to ``_statistics_methods``. |
| | |
| | The registered method can be called by ``history_buffer.statistics`` |
| | with corresponding method name and arguments. |
| | |
| | Examples: |
| | >>> @HistoryBuffer.register_statistics |
| | >>> def weighted_mean(self, window_size, weight): |
| | >>> assert len(weight) == window_size |
| | >>> return (self._log_history[-window_size:] * |
| | >>> np.array(weight)).sum() / \ |
| | >>> self._count_history[-window_size:] |
| | |
| | >>> log_buffer = HistoryBuffer([1, 2], [1, 1]) |
| | >>> log_buffer.statistics('weighted_mean', 2, [2, 1]) |
| | 2 |
| | |
| | Args: |
| | method (Callable): Custom statistics method. |
| | Returns: |
| | Callable: Original custom statistics method. |
| | """ |
| | method_name = method.__name__ |
| | assert method_name not in cls._statistics_methods, \ |
| | 'method_name cannot be registered twice!' |
| | cls._statistics_methods[method_name] = method |
| | return method |
| |
|
| | def statistics(self, method_name: str, *arg, **kwargs) -> Any: |
| | """Access statistics method by name. |
| | |
| | Args: |
| | method_name (str): Name of method. |
| | |
| | Returns: |
| | Any: Depends on corresponding method. |
| | """ |
| | if method_name not in self._statistics_methods: |
| | raise KeyError(f'{method_name} has not been registered in ' |
| | 'HistoryBuffer._statistics_methods') |
| | method = self._statistics_methods[method_name] |
| | |
| | return method(self, *arg, **kwargs) |
| |
|
| | def mean(self, window_size: Optional[int] = None) -> np.ndarray: |
| | """Return the mean of the latest ``window_size`` values in log |
| | histories. |
| | |
| | If ``window_size is None`` or ``window_size > len(self._log_history)``, |
| | return the global mean value of history logs. |
| | |
| | Args: |
| | window_size (int, optional): Size of statistics window. |
| | Returns: |
| | np.ndarray: Mean value within the window. |
| | """ |
| | if window_size is not None: |
| | assert isinstance(window_size, int), \ |
| | 'The type of window size should be int, but got ' \ |
| | f'{type(window_size)}' |
| | else: |
| | window_size = len(self._log_history) |
| | logs_sum = self._log_history[-window_size:].sum() |
| | counts_sum = self._count_history[-window_size:].sum() |
| | return logs_sum / counts_sum |
| |
|
| | def max(self, window_size: Optional[int] = None) -> np.ndarray: |
| | """Return the maximum value of the latest ``window_size`` values in log |
| | histories. |
| | |
| | If ``window_size is None`` or ``window_size > len(self._log_history)``, |
| | return the global maximum value of history logs. |
| | |
| | Args: |
| | window_size (int, optional): Size of statistics window. |
| | Returns: |
| | np.ndarray: The maximum value within the window. |
| | """ |
| | if window_size is not None: |
| | assert isinstance(window_size, int), \ |
| | 'The type of window size should be int, but got ' \ |
| | f'{type(window_size)}' |
| | else: |
| | window_size = len(self._log_history) |
| | return self._log_history[-window_size:].max() |
| |
|
| | def min(self, window_size: Optional[int] = None) -> np.ndarray: |
| | """Return the minimum value of the latest ``window_size`` values in log |
| | histories. |
| | |
| | If ``window_size is None`` or ``window_size > len(self._log_history)``, |
| | return the global minimum value of history logs. |
| | |
| | Args: |
| | window_size (int, optional): Size of statistics window. |
| | Returns: |
| | np.ndarray: The minimum value within the window. |
| | """ |
| | if window_size is not None: |
| | assert isinstance(window_size, int), \ |
| | 'The type of window size should be int, but got ' \ |
| | f'{type(window_size)}' |
| | else: |
| | window_size = len(self._log_history) |
| | return self._log_history[-window_size:].min() |
| |
|
| | def current(self) -> np.ndarray: |
| | """Return the recently updated values in log histories. |
| | |
| | Returns: |
| | np.ndarray: Recently updated values in log histories. |
| | """ |
| | if len(self._log_history) == 0: |
| | raise ValueError('HistoryBuffer._log_history is an empty array! ' |
| | 'please call update first') |
| | return self._log_history[-1] |
| |
|
| | def __getstate__(self) -> dict: |
| | """Make ``_statistics_methods`` can be resumed. |
| | |
| | Returns: |
| | dict: State dict including statistics_methods. |
| | """ |
| | self.__dict__.update(statistics_methods=self._statistics_methods) |
| | return self.__dict__ |
| |
|
| | def __setstate__(self, state): |
| | """Try to load ``_statistics_methods`` from state. |
| | |
| | Args: |
| | state (dict): State dict. |
| | """ |
| | statistics_methods = state.pop('statistics_methods', {}) |
| | self._set_default_statistics() |
| | self._statistics_methods.update(statistics_methods) |
| | self.__dict__.update(state) |
| |
|
