Function sync_batch_norm

imperative/python/megengine/functional/nn.py:1475–1609

r"""Applies synchronized batch normalization to the input. Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information. Args: inp: input tensor. running_mean: tensor to store running mean. running_var: tensor to store running variance. we

(
    inp: Tensor,
    running_mean: Tensor,
    running_var: Tensor,
    weight: Optional[Tensor] = None,
    bias: Optional[Tensor] = None,
    training: bool = False,
    momentum: Union[float, Tensor] = 0.9,
    eps: float = 1e-5,
    eps_mode="additive",
    group=WORLD,
)

def sync_batch_norm(
    inp: Tensor,
    running_mean: Tensor,
    running_var: Tensor,
    weight: Optional[Tensor] = None,
    bias: Optional[Tensor] = None,
    training: bool = False,
    momentum: Union[float, Tensor] = 0.9,
    eps: float = 1e-5,
    eps_mode="additive",
    group=WORLD,
) -> Tensor:
    r"""Applies synchronized batch normalization to the input.

    Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information.

    Args:
        inp: input tensor.
        running_mean: tensor to store running mean.
        running_var: tensor to store running variance.
        weight: scaling tensor in the learnable affine parameters.
            See :math:`\gamma` in :class:`~.BatchNorm2d`.
        bias: bias tensor in the learnable affine parameters.
            See :math:`\beta` in :class:`~.BatchNorm2d`.
        training: a boolean value to indicate whether batch norm is performed
            in traning mode. Default: False
        momentum: value used for the ``running_mean`` and ``running_var``
            computation. Default: 0.9
        eps: a value added to the denominator for numerical stability.
            Default: 1e-5
        eps_mode: mode of calculation for eps, "max" or "additive".
            Default: "additive"
        group: communication group, caculate mean and variance between this group.
            Default: :obj:`~megengine.distributed.WORLD`
    """
    _eps_mode = eps_mode.lower()
    assert _eps_mode in {"max", "additive"}, "unknown eps_mode: {}".format(eps_mode)
    if _eps_mode == "additive" and not (is_distributed() and training):
        return batch_norm(
            inp,
            running_mean,
            running_var,
            weight,
            bias,
            training=training,
            momentum=momentum,
            eps=eps,
        )
    if amp._enabled:
        inp, weight, bias, running_mean, running_var = cast_tensors(
            inp, weight, bias, running_mean, running_var, promote=True
        )

    _channels = make_shape_tuple(inp.shape)[1]
    _ndim = inp.ndim
    _device = inp.device
    _dtype = inp.dtype