Class SyncBatchNorm

imperative/python/megengine/module/batchnorm.py:117–214

r"""Applies Synchronized Batch Normalization for distributed training. Args: num_features: usually :math:`C` from an input of shape :math:`(N, C, H, W)` or the highest ranked dimension of an input less than 4D. eps: a value added to the denominator for nu

class SyncBatchNorm(_BatchNorm):
    r"""Applies Synchronized Batch Normalization for distributed training.

    Args:
        num_features: usually :math:`C` from an input of shape
            :math:`(N, C, H, W)` or the highest ranked dimension of an input
            less than 4D.
        eps: a value added to the denominator for numerical stability.
            Default: 1e-5
        momentum: the value used for the ``running_mean`` and ``running_var`` computation.
            Default: 0.9
        affine: a boolean value that when set to True, this module has
            learnable affine parameters. Default: True
        track_running_stats: when set to True, this module tracks the
            running mean and variance. When set to False, this module does not
            track such statistics and always uses batch statistics in both training
            and eval modes. Default: True
        freeze: when set to True, this module does not update the
            running mean and variance, and uses the running mean and variance instead of
            the batch mean and batch variance to normalize the input. The parameter takes effect
            only when the module is initilized with track_running_stats as True.
            Default: False
        group: communication group, caculate mean and variance between this group.
            Default: :obj:`~.distributed.WORLD`
    """

    def __init__(
        self,
        num_features,
        eps=1e-5,
        momentum=0.9,
        affine=True,
        track_running_stats=True,
        freeze=False,
        group: Optional[Group] = WORLD,
        **kwargs
    ) -> None:
        super().__init__(
            num_features, eps, momentum, affine, track_running_stats, freeze, **kwargs
        )
        self.group = group

    def _check_input_ndim(self, inp):
        if len(inp.shape) not in {2, 3, 4}:
            raise ValueError(
                "expected 2D, 3D or 4D input (got {}D input)".format(len(inp.shape))
            )

    def forward(self, inp):
        self._check_input_ndim(inp)

        inp_shape = inp.shape
        _ndims = len(inp_shape)
        if _ndims != 4:
            new_shape = Tensor([1, 1, 1, 1], device=inp.device)
            origin_shape = inp_shape
            if _ndims == 2:
                new_shape[:2] = origin_shape[:2]