batch_agg.py

import operator
from enum import Enum
from functools import reduce
from typing import *

import numpy as np

from .stage import StageType
from .utils import ALL, NOT_SET

__all__ = [
    'BatchAggregationMode',
    'BatchAggregator', 'BatchAggregatorDict',
]


class BatchAggregationMode(str, Enum):

    CONCAT = 'CONCAT'
    """To concat the batch arrays along specified axis."""

    SUM = 'SUM'
    """To sum the batch arrays along specified axis."""

    AVERAGE = 'AVERAGE'
    """To average the batch arrays along specified axis."""


class BatchAggregator(object):
    """
    Class to aggregate batch arrays.

    >>> agg = BatchAggregator(BatchAggregationMode.CONCAT)
    >>> agg
    BatchAggregator(mode=CONCAT, axis=0)
    >>> agg.add(np.array([1, 2, 3, 4]))
    >>> agg.add(np.array([5, 6]))
    >>> agg.get()
    array([1, 2, 3, 4, 5, 6])

    >>> agg = BatchAggregator(BatchAggregationMode.AVERAGE)
    >>> agg
    BatchAggregator(mode=AVERAGE, axis=None)
    >>> agg.add(np.array([1, 2, 3, 4]))
    >>> agg.add(np.array([5, 6]))
    >>> agg.get()
    3.5

    >>> agg = BatchAggregator(BatchAggregationMode.SUM)
    >>> agg
    BatchAggregator(mode=SUM, axis=None)
    >>> agg.add(np.array([1, 2, 3, 4]))
    >>> agg.add(np.array([5, 6]))
    >>> agg.get()
    21
    """

    mode: BatchAggregationMode
    axis: Union[int, Tuple[int, ...]]

    def __init__(self,
                 mode: Union[str, BatchAggregationMode],
                 axis: Optional[Union[int, Tuple[int, ...], List[int]]] = NOT_SET):
        """
        Construct a new :class:`BatchAggregator`.

        Args:
            mode: Aggregation mode.
            axis: The axis to aggregate.  Defaults to `0` for `CONCAT` mode,
                while :obj:`None` for `SUM` and `AVERAGE` mode.
        """
        mode = BatchAggregationMode(mode)

        if axis is NOT_SET:
            axis = 0 if mode == BatchAggregationMode.CONCAT else None
        if mode == BatchAggregationMode.CONCAT:
            if not isinstance(axis, int):
                raise TypeError('`axis` must be a int when `mode` is CONCAT.')
        if axis is not None:
            if hasattr(axis, '__iter__'):
                axis = tuple(int(v) for v in axis)
                if len(axis) == 1:
                    axis = axis[0]
            else:
                axis = int(axis)

        self.mode = mode
        self.axis = axis
        self._buf = None
        self._weight_sum = 0.

    def __repr__(self):
        return f'{self.__class__.__qualname__}' \
               f'(mode={self.mode.value}, axis={self.axis})'

    def get(self) -> Optional[np.ndarray]:
        """
        Get the aggregation result.

        Returns:
            The result, or :obj:`None` if no value has been collected.
        """
        if self._buf is not None:
            if self.mode == BatchAggregationMode.CONCAT:
                return np.concatenate(self._buf, axis=self.axis)
            else:
                return self._buf

    def add(self,
            values: np.ndarray,
            weight: Optional[float] = 1.):
        """
        Add a batch array to the aggregator.

        Args:
            values: The batch array.
            weight: The batch weight, used only in `AVERAGE` mode.
        """
        # CONCAT: append the values to the buf
        if self.mode == BatchAggregationMode.CONCAT:
            if self._buf is None:
                self._buf = []
            self._buf.append(values)

        # SUM
        elif self.mode == BatchAggregationMode.SUM:
            batch_sum = np.sum(values, axis=self.axis)
            if self._buf is None:
                self._buf = batch_sum
            else:
                self._buf += batch_sum

        # AVERAGE: maintain the `total_weight` state and update the buf
        else:
            # infer the batch size and weight
            batch_shape = np.shape(values)
            if self.axis is None:
                batch_size = float(reduce(operator.mul, np.shape(values), 1.))
            elif isinstance(self.axis, tuple):
                batch_size = 1.
                for a in self.axis:
                    batch_size *= batch_shape[a]
            else:
                batch_size = batch_shape[self.axis]
            batch_weight = weight * batch_size

            # do update the weight
            self._weight_sum += batch_weight
            r1 = weight / self._weight_sum
            batch_sum = np.sum(values, axis=self.axis)
            if self._buf is None:
                self._buf = r1 * batch_sum
            else:
                r2 = batch_weight / self._weight_sum
                self._buf += r1 * batch_sum - r2 * self._buf


class BatchAggregatorDict(Mapping[str, BatchAggregator]):
    """
    Maintain a dict of :class:`BatchAggregator` instances, maybe with
    a default factory to construct :class:`BatchAggregator` instance
    for new keys.

    >>> agg_dict = BatchAggregatorDict.new()
    >>> agg_dict['acc'].add(np.array([0.75, 0.875]))
    >>> agg_dict['loss'].add(np.array([0.125, 0.2]))
    >>> len(agg_dict)
    2
    >>> list(agg_dict)
    ['acc', 'loss']
    >>> agg_dict['acc'].get()
    0.8125
    >>> agg_dict['loss'].get()
    0.1625
    """

    @staticmethod
    def new(metrics: Union[Sequence[str], type(ALL)] = ALL,
            outputs: Union[Sequence[str], type(ALL)] = (),
            aggregators: Optional[Mapping[str, BatchAggregator]] = None,
            excludes: Sequence[str] = (),
            stage_type: Optional[StageType] = None) -> 'BatchAggregatorDict':
        """
        Construct a new :class:`BatchAggregatorDict` according to the field
        settings `metrics`, `outputs` and `aggregators`.

        Args:
            metrics: The names of the batch arrays, which should be aggregated
                by ``BatchAggregator('AVERAGE', axis=None)``.  :obj:`ALL`
                indicates that an array is by default a metric if it is neither
                specified in `outputs` nor in `aggregator`.
            outputs: The names of the batch arrays, which should be aggregated
                by ``BatchAggregator('CONCAT', axis=0)``.  :obj:`ALL`
                indicates that an array is by default an output if it is neither
                specified in `outputs` nor in `aggregator`.
            aggregators: The dict of names and their corresponding aggregators.
            excludes: The names to exclude.  If a name is excluded, no
                aggregator will be designated to this name, i.e., ``get(name)``
                returns None, and ``__getitem__(name)`` raises `KeyError`.
            stage_type: If specified, will add stage metric prefix to the keys
                of `metrics`, `outputs` and `aggregators`.

        Returns:
            The aggregator dict.

        Notes:
            :obj:`ALL` could be specified to at most one of `metrics`
            and `outputs`.  The argument `aggregators` has higher priority
            than `outputs`, and so does `outputs` have higher priority than
            `metrics`.  That is to say, if a name is specified in both
            `aggregators` and `outputs`, then the aggregator specified in
            `aggregators` will be chosen; this is also true if a name is
            specified in both `outputs` and `metrics`.
        """
        # the aggregator factories
        average_aggregator_factory = lambda: \
            BatchAggregator(mode=BatchAggregationMode.AVERAGE, axis=None)
        concat_aggregator_factory = lambda: \
            BatchAggregator(mode=BatchAggregationMode.CONCAT, axis=0)

        # determine the default factory
        if metrics == ALL and outputs == ALL:
            raise ValueError('Only one of `metrics` and `outputs` can be '
                             '`ALL`.')
        elif metrics == ALL:
            default_factory = average_aggregator_factory
        elif outputs == ALL:
            default_factory = concat_aggregator_factory
        else:
            default_factory = None

        # build the aggregator instances
        agg_dict = {}
        if metrics != ALL and metrics:
            for key in metrics:
                if stage_type is not None:
                    key = stage_type.add_metric_prefix(key)
                agg_dict[key] = average_aggregator_factory()
        if outputs != ALL and outputs:
            for key in outputs:
                if stage_type is not None:
                    key = stage_type.add_metric_prefix(key)
                agg_dict[key] = concat_aggregator_factory()
        if aggregators:
            for key, agg in aggregators.items():
                if stage_type is not None:
                    key = stage_type.add_metric_prefix(key)
                agg_dict[key] = agg

        # build the excludes names
        if excludes and stage_type is not None:
            excludes = [stage_type.add_metric_prefix(n) for n in excludes]

        # now construct the `BatchAggregatorDict` instance
        return BatchAggregatorDict(
            agg_dict, excludes=excludes, default_factory=default_factory)

    def __init__(self,
                 aggregators: Mapping[str, BatchAggregator],
                 excludes: Sequence[str] = (),
                 default_factory: Optional[
                     Callable[[], BatchAggregator]] = None):
        """
        Construct a new :class:`BatchAggregatorDict`.

        Args:
            aggregators: The mapping from names to aggregators.
            excludes: The names to exclude from this dict.  If a name is
                excluded, no aggregator will be designated to this name,
                i.e., ``get(name)`` returns None, and ``__getitem__(name)``
                raises :class:`KeyError`.
            default_factory: The default factory, which is used to create
                new :class:`BatchAggregator` instances if the aggregator
                to a requested name does not exist.  If not specified,
                accessing non-existing name will raise an error.
        """
        self._aggregators = {}
        self._excludes = set(excludes or ())
        self._default_factory = default_factory

        for key in aggregators:
            if key not in self._excludes:
                agg = aggregators[key]
                if not isinstance(agg, BatchAggregator):
                    raise TypeError(f'Item {key!r} is not an instance of '
                                    f'{BatchAggregator.__qualname__}: '
                                    f'{agg!r}')
                self._aggregators[key] = agg

    def get(self, item: str, default: Any = None) -> Optional[BatchAggregator]:
        if item not in self._excludes:
            if item not in self._aggregators:
                if self._default_factory is not None:
                    self._aggregators[item] = self._default_factory()
                else:
                    return default
            return self._aggregators[item]

    def __getitem__(self, item: str) -> BatchAggregator:
        ret = self.get(item)
        if ret is None:
            raise KeyError(item)
        return ret

    def __len__(self) -> int:
        return len(self._aggregators)

    def __iter__(self) -> Iterator[str]:
        return iter(self._aggregators)