diff --git a/python_module/megengine/functional/nn.py b/python_module/megengine/functional/nn.py index d781fef1a14320d00f97d3a13a61daec866e60dc..ec9f670d9cf51bc0ef78b90bd840d3571c40813d 100644 --- a/python_module/megengine/functional/nn.py +++ b/python_module/megengine/functional/nn.py @@ -280,27 +280,20 @@ def batch_norm2d( ) -> Tensor: """Applies batch normalization to the input. - :type inp: Tensor - :param inp: The input tensor. - :type num_features: int - :param num_features: usually the :math:`C` from an input of size - :math:`(N, C, H, W)` or the highest ranked dimension of an input with - less than 4D. - :type eps: float + :param inp: input tensor. + :param running_mean: tensor to store running mean. + :param running_var: tensor to store running variance. + :param weight: scaling tensor in the learnable affine parameters. + See :math:`\gamma` in :class:`~.BatchNorm2d` + :param bias: bias tensor in the learnable affine parameters. + See :math:`\beta` in :class:`~.BatchNorm2d` + :param training: a boolean value to indicate whether batch norm is performed + in traning mode. Default: ``False`` + :param momentum: the value used for the ``running_mean`` and ``running_var`` + computation. + Default: 0.9 :param eps: a value added to the denominator for numerical stability. Default: 1e-5. - :type momentum: float - :param momentum: the value used for the `running_mean` and `running_var` - computation. - Default: 0.1 - :type affine: bool - :param affine: a boolean value that when set to ``True``, this module has - learnable affine parameters. Default: ``True`` - :type track_running_stats: bool - :param 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``. Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information. """