Table of Contents
Shortcuts

# Source code for torch.nn.modules.batchnorm

    from typing import Optional, Any

    import torch
    from torch import Tensor
    from torch.nn.parameter import Parameter, UninitializedParameter, UninitializedBuffer

    from .. import functional as F
    from .. import init
    from ._functions import SyncBatchNorm as sync_batch_norm
    from .lazy import LazyModuleMixin
    from .module import Module

    __all__ = ['BatchNorm1d', 'LazyBatchNorm1d', 'BatchNorm2d', 'LazyBatchNorm2d', 'BatchNorm3d',
               'LazyBatchNorm3d', 'SyncBatchNorm']

    class _NormBase(Module):
        """Common base of _InstanceNorm and _BatchNorm"""

        _version = 2
        __constants__ = ["track_running_stats", "momentum", "eps", "num_features", "affine"]
        num_features: int
        eps: float
        momentum: float
        affine: bool
        track_running_stats: bool
        # WARNING: weight and bias purposely not defined here.
        # See https://github.com/pytorch/pytorch/issues/39670

        def __init__(
            self,
            num_features: int,
            eps: float = 1e-5,
            momentum: float = 0.1,
            affine: bool = True,
            track_running_stats: bool = True,
            device=None,
            dtype=None
        ) -> None:
            factory_kwargs = {'device': device, 'dtype': dtype}
            super().__init__()
            self.num_features = num_features
            self.eps = eps
            self.momentum = momentum
            self.affine = affine
            self.track_running_stats = track_running_stats
            if self.affine:
                self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
                self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
            else:
                self.register_parameter("weight", None)
                self.register_parameter("bias", None)
            if self.track_running_stats:
                self.register_buffer('running_mean', torch.zeros(num_features, **factory_kwargs))
                self.register_buffer('running_var', torch.ones(num_features, **factory_kwargs))
                self.running_mean: Optional[Tensor]
                self.running_var: Optional[Tensor]
                self.register_buffer('num_batches_tracked',
                                     torch.tensor(0, dtype=torch.long,
                                                  **{k: v for k, v in factory_kwargs.items() if k != 'dtype'}))
                self.num_batches_tracked: Optional[Tensor]
            else:
                self.register_buffer("running_mean", None)
                self.register_buffer("running_var", None)
                self.register_buffer("num_batches_tracked", None)
            self.reset_parameters()

        def reset_running_stats(self) -> None:
            if self.track_running_stats:
                # running_mean/running_var/num_batches... are registered at runtime depending
                # if self.track_running_stats is on
                self.running_mean.zero_()  # type: ignore[union-attr]
                self.running_var.fill_(1)  # type: ignore[union-attr]
                self.num_batches_tracked.zero_()  # type: ignore[union-attr,operator]

        def reset_parameters(self) -> None:
            self.reset_running_stats()
            if self.affine:
                init.ones_(self.weight)
                init.zeros_(self.bias)

        def _check_input_dim(self, input):
            raise NotImplementedError

        def extra_repr(self):
            return (
                "{num_features}, eps={eps}, momentum={momentum}, affine={affine}, "
                "track_running_stats={track_running_stats}".format(**self.__dict__)
            )

        def _load_from_state_dict(
            self,
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        ):
            version = local_metadata.get("version", None)

            if (version is None or version < 2) and self.track_running_stats:
                # at version 2: added num_batches_tracked buffer
                #               this should have a default value of 0
                num_batches_tracked_key = prefix + "num_batches_tracked"
                if num_batches_tracked_key not in state_dict:
                    state_dict[num_batches_tracked_key] = torch.tensor(0, dtype=torch.long)

            super()._load_from_state_dict(
                state_dict,
                prefix,
                local_metadata,
                strict,
                missing_keys,
                unexpected_keys,
                error_msgs,
            )

    class _BatchNorm(_NormBase):
        def __init__(
            self,
            num_features: int,
            eps: float = 1e-5,
            momentum: float = 0.1,
            affine: bool = True,
            track_running_stats: bool = True,
            device=None,
            dtype=None
        ) -> None:
            factory_kwargs = {'device': device, 'dtype': dtype}
            super().__init__(
                num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
            )

        def forward(self, input: Tensor) -> Tensor:
            self._check_input_dim(input)

            # exponential_average_factor is set to self.momentum
            # (when it is available) only so that it gets updated
            # in ONNX graph when this node is exported to ONNX.
            if self.momentum is None:
                exponential_average_factor = 0.0
            else:
                exponential_average_factor = self.momentum

            if self.training and self.track_running_stats:
                # TODO: if statement only here to tell the jit to skip emitting this when it is None
                if self.num_batches_tracked is not None:  # type: ignore[has-type]
                    self.num_batches_tracked.add_(1)  # type: ignore[has-type]
                    if self.momentum is None:  # use cumulative moving average
                        exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                    else:  # use exponential moving average
                        exponential_average_factor = self.momentum

            r"""
            Decide whether the mini-batch stats should be used for normalization rather than the buffers.
            Mini-batch stats are used in training mode, and in eval mode when buffers are None.
            """
            if self.training:
                bn_training = True
            else:
                bn_training = (self.running_mean is None) and (self.running_var is None)

            r"""
            Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
            passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
            used for normalization (i.e. in eval mode when buffers are not None).
            """
            return F.batch_norm(
                input,
                # If buffers are not to be tracked, ensure that they won't be updated
                self.running_mean
                if not self.training or self.track_running_stats
                else None,
                self.running_var if not self.training or self.track_running_stats else None,
                self.weight,
                self.bias,
                bn_training,
                exponential_average_factor,
                self.eps,
            )

    class _LazyNormBase(LazyModuleMixin, _NormBase):

        weight: UninitializedParameter  # type: ignore[assignment]
        bias: UninitializedParameter  # type: ignore[assignment]

        def __init__(self, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True,
                     device=None, dtype=None) -> None:
            factory_kwargs = {'device': device, 'dtype': dtype}
            super().__init__(
                # affine and track_running_stats are hardcoded to False to
                # avoid creating tensors that will soon be overwritten.
                0,
                eps,
                momentum,
                False,
                False,
                **factory_kwargs,
            )
            self.affine = affine
            self.track_running_stats = track_running_stats
            if self.affine:
                self.weight = UninitializedParameter(**factory_kwargs)
                self.bias = UninitializedParameter(**factory_kwargs)
            if self.track_running_stats:
                self.running_mean = UninitializedBuffer(**factory_kwargs)
                self.running_var = UninitializedBuffer(**factory_kwargs)
                self.num_batches_tracked = torch.tensor(
                    0, dtype=torch.long, **{k: v for k, v in factory_kwargs.items() if k != 'dtype'})

        def reset_parameters(self) -> None:
            if not self.has_uninitialized_params() and self.num_features != 0:
                super().reset_parameters()

        def initialize_parameters(self, input) -> None:  # type: ignore[override]
            if self.has_uninitialized_params():
                self.num_features = input.shape[1]
                if self.affine:
                    assert isinstance(self.weight, UninitializedParameter)
                    assert isinstance(self.bias, UninitializedParameter)
                    self.weight.materialize((self.num_features,))
                    self.bias.materialize((self.num_features,))
                if self.track_running_stats:
                    self.running_mean.materialize((self.num_features,))  # type:ignore[union-attr]
                    self.running_var.materialize((self.num_features,))  # type:ignore[union-attr]
                self.reset_parameters()

    class BatchNorm1d(_BatchNorm):
        r"""Applies Batch Normalization over a 2D or 3D input as described in the paper
        `Batch Normalization: Accelerating Deep Network Training by Reducing
        Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .

        .. math::

            y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

        The mean and standard-deviation are calculated per-dimension over
        the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
        of size `C` (where `C` is the number of features or channels of the input). By default, the
        elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
        At train time in the forward pass, the standard-deviation is calculated via the biased estimator,
        equivalent to ``torch.var(input, unbiased=False)``. However, the value stored in the
        moving average of the standard-deviation is calculated via the unbiased  estimator, equivalent to
        ``torch.var(input, unbiased=True)``.

        Also by default, during training this layer keeps running estimates of its
        computed mean and variance, which are then used for normalization during
        evaluation. The running estimates are kept with a default :attr:`momentum`
        of 0.1.

        If :attr:`track_running_stats` is set to ``False``, this layer then does not
        keep running estimates, and batch statistics are instead used during
        evaluation time as well.

        .. note::
            This :attr:`momentum` argument is different from one used in optimizer
            classes and the conventional notion of momentum. Mathematically, the
            update rule for running statistics here is
            :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
            where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
            new observed value.

        Because the Batch Normalization is done over the `C` dimension, computing statistics
        on `(N, L)` slices, it's common terminology to call this Temporal Batch Normalization.

        Args:
            num_features: number of features or channels :math:`C` of the input
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``

        Shape:
            - Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size,
              :math:`C` is the number of features or channels, and :math:`L` is the sequence length
            - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)

        Examples::

            >>> # With Learnable Parameters
            >>> m = nn.BatchNorm1d(100)
            >>> # Without Learnable Parameters
            >>> m = nn.BatchNorm1d(100, affine=False)
            >>> input = torch.randn(20, 100)
            >>> output = m(input)
        """

        def _check_input_dim(self, input):
            if input.dim() != 2 and input.dim() != 3:
                raise ValueError(
                    f"expected 2D or 3D input (got {input.dim()}D input)"
                )

    class LazyBatchNorm1d(_LazyNormBase, _BatchNorm):
        r"""A :class:`torch.nn.BatchNorm1d` module with lazy initialization of
        the ``num_features`` argument of the :class:`BatchNorm1d` that is inferred
        from the ``input.size(1)``.
        The attributes that will be lazily initialized are `weight`, `bias`,
        `running_mean` and `running_var`.

        Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
        on lazy modules and their limitations.

        Args:
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``
        """

        cls_to_become = BatchNorm1d  # type: ignore[assignment]

        def _check_input_dim(self, input):
            if input.dim() != 2 and input.dim() != 3:
                raise ValueError(
                    f"expected 2D or 3D input (got {input.dim()}D input)"
                )

    class BatchNorm2d(_BatchNorm):
        r"""Applies Batch Normalization over a 4D input (a mini-batch of 2D inputs
        with additional channel dimension) as described in the paper
        `Batch Normalization: Accelerating Deep Network Training by Reducing
        Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .

        .. math::

            y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

        The mean and standard-deviation are calculated per-dimension over
        the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
        of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
        to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
        standard-deviation is calculated via the biased estimator, equivalent to
        ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the
        standard-deviation is calculated via the unbiased  estimator, equivalent to
        ``torch.var(input, unbiased=True)``.

        Also by default, during training this layer keeps running estimates of its
        computed mean and variance, which are then used for normalization during
        evaluation. The running estimates are kept with a default :attr:`momentum`
        of 0.1.

        If :attr:`track_running_stats` is set to ``False``, this layer then does not
        keep running estimates, and batch statistics are instead used during
        evaluation time as well.

        .. note::
            This :attr:`momentum` argument is different from one used in optimizer
            classes and the conventional notion of momentum. Mathematically, the
            update rule for running statistics here is
            :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
            where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
            new observed value.

        Because the Batch Normalization is done over the `C` dimension, computing statistics
        on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization.

        Args:
            num_features: :math:`C` from an expected input of size
                :math:`(N, C, H, W)`
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``

        Shape:
            - Input: :math:`(N, C, H, W)`
            - Output: :math:`(N, C, H, W)` (same shape as input)

        Examples::

            >>> # With Learnable Parameters
            >>> m = nn.BatchNorm2d(100)
            >>> # Without Learnable Parameters
            >>> m = nn.BatchNorm2d(100, affine=False)
            >>> input = torch.randn(20, 100, 35, 45)
            >>> output = m(input)
        """

        def _check_input_dim(self, input):
            if input.dim() != 4:
                raise ValueError(f"expected 4D input (got {input.dim()}D input)")

    class LazyBatchNorm2d(_LazyNormBase, _BatchNorm):
        r"""A :class:`torch.nn.BatchNorm2d` module with lazy initialization of
        the ``num_features`` argument of the :class:`BatchNorm2d` that is inferred
        from the ``input.size(1)``.
        The attributes that will be lazily initialized are `weight`, `bias`,
        `running_mean` and `running_var`.

        Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
        on lazy modules and their limitations.

        Args:
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``
        """

        cls_to_become = BatchNorm2d  # type: ignore[assignment]

        def _check_input_dim(self, input):
            if input.dim() != 4:
                raise ValueError(f"expected 4D input (got {input.dim()}D input)")

    class BatchNorm3d(_BatchNorm):
        r"""Applies Batch Normalization over a 5D input (a mini-batch of 3D inputs
        with additional channel dimension) as described in the paper
        `Batch Normalization: Accelerating Deep Network Training by Reducing
        Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .

        .. math::

            y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

        The mean and standard-deviation are calculated per-dimension over
        the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
        of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
        to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
        standard-deviation is calculated via the biased estimator, equivalent to
        ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the
        standard-deviation is calculated via the unbiased  estimator, equivalent to
        ``torch.var(input, unbiased=True)``.

        Also by default, during training this layer keeps running estimates of its
        computed mean and variance, which are then used for normalization during
        evaluation. The running estimates are kept with a default :attr:`momentum`
        of 0.1.

        If :attr:`track_running_stats` is set to ``False``, this layer then does not
        keep running estimates, and batch statistics are instead used during
        evaluation time as well.

        .. note::
            This :attr:`momentum` argument is different from one used in optimizer
            classes and the conventional notion of momentum. Mathematically, the
            update rule for running statistics here is
            :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
            where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
            new observed value.

        Because the Batch Normalization is done over the `C` dimension, computing statistics
        on `(N, D, H, W)` slices, it's common terminology to call this Volumetric Batch Normalization
        or Spatio-temporal Batch Normalization.

        Args:
            num_features: :math:`C` from an expected input of size
                :math:`(N, C, D, H, W)`
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``

        Shape:
            - Input: :math:`(N, C, D, H, W)`
            - Output: :math:`(N, C, D, H, W)` (same shape as input)

        Examples::

            >>> # With Learnable Parameters
            >>> m = nn.BatchNorm3d(100)
            >>> # Without Learnable Parameters
            >>> m = nn.BatchNorm3d(100, affine=False)
            >>> input = torch.randn(20, 100, 35, 45, 10)
            >>> output = m(input)
        """

        def _check_input_dim(self, input):
            if input.dim() != 5:
                raise ValueError(f"expected 5D input (got {input.dim()}D input)")

    class LazyBatchNorm3d(_LazyNormBase, _BatchNorm):
        r"""A :class:`torch.nn.BatchNorm3d` module with lazy initialization of
        the ``num_features`` argument of the :class:`BatchNorm3d` that is inferred
        from the ``input.size(1)``.
        The attributes that will be lazily initialized are `weight`, `bias`,
        `running_mean` and `running_var`.

        Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
        on lazy modules and their limitations.

        Args:
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``
        """

        cls_to_become = BatchNorm3d  # type: ignore[assignment]

        def _check_input_dim(self, input):
            if input.dim() != 5:
                raise ValueError(f"expected 5D input (got {input.dim()}D input)")

    class SyncBatchNorm(_BatchNorm):
        r"""Applies Batch Normalization over a N-Dimensional input (a mini-batch of [N-2]D inputs
        with additional channel dimension) as described in the paper
        `Batch Normalization: Accelerating Deep Network Training by Reducing
        Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .

        .. math::

            y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

        The mean and standard-deviation are calculated per-dimension over all
        mini-batches of the same process groups. :math:`\gamma` and :math:`\beta`
        are learnable parameter vectors of size `C` (where `C` is the input size).
        By default, the elements of :math:`\gamma` are sampled from
        :math:`\mathcal{U}(0, 1)` and the elements of :math:`\beta` are set to 0.
        The standard-deviation is calculated via the biased estimator, equivalent to
        `torch.var(input, unbiased=False)`.

        Also by default, during training this layer keeps running estimates of its
        computed mean and variance, which are then used for normalization during
        evaluation. The running estimates are kept with a default :attr:`momentum`
        of 0.1.

        If :attr:`track_running_stats` is set to ``False``, this layer then does not
        keep running estimates, and batch statistics are instead used during
        evaluation time as well.

        .. note::
            This :attr:`momentum` argument is different from one used in optimizer
            classes and the conventional notion of momentum. Mathematically, the
            update rule for running statistics here is
            :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
            where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
            new observed value.

        Because the Batch Normalization is done for each channel in the ``C`` dimension, computing
        statistics on ``(N, +)`` slices, it's common terminology to call this Volumetric Batch
        Normalization or Spatio-temporal Batch Normalization.

        Currently :class:`SyncBatchNorm` only supports
        :class:`~torch.nn.DistributedDataParallel` (DDP) with single GPU per process. Use
        :meth:`torch.nn.SyncBatchNorm.convert_sync_batchnorm()` to convert
        :attr:`BatchNorm*D` layer to :class:`SyncBatchNorm` before wrapping
        Network with DDP.

        Args:
            num_features: :math:`C` from an expected input of size
                :math:`(N, C, +)`
            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. Can be set to ``None`` for cumulative moving average
                (i.e. simple average). Default: 0.1
            affine: a boolean value that when set to ``True``, this module has
                learnable affine parameters. Default: ``True``
            track_running_stats: a boolean value that when set to ``True``, this
                module tracks the running mean and variance, and when set to ``False``,
                this module does not track such statistics, and initializes statistics
                buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
                When these buffers are ``None``, this module always uses batch statistics.
                in both training and eval modes. Default: ``True``
            process_group: synchronization of stats happen within each process group
                individually. Default behavior is synchronization across the whole
                world

        Shape:
            - Input: :math:`(N, C, +)`
            - Output: :math:`(N, C, +)` (same shape as input)

        .. note::
            Synchronization of batchnorm statistics occurs only while training, i.e.
            synchronization is disabled when ``model.eval()`` is set or if
            ``self.training`` is otherwise ``False``.

        Examples::

            >>> # xdoctest: +SKIP
            >>> # With Learnable Parameters
            >>> m = nn.SyncBatchNorm(100)
            >>> # creating process group (optional)
            >>> # ranks is a list of int identifying rank ids.
            >>> ranks = list(range(8))
            >>> r1, r2 = ranks[:4], ranks[4:]
            >>> # Note: every rank calls into new_group for every
            >>> # process group created, even if that rank is not
            >>> # part of the group.
            >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
            >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
            >>> # Without Learnable Parameters
            >>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)
            >>> input = torch.randn(20, 100, 35, 45, 10)
            >>> output = m(input)

            >>> # network is nn.BatchNorm layer
            >>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)
            >>> # only single gpu per process is currently supported
            >>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
            >>>                         sync_bn_network,
            >>>                         device_ids=[args.local_rank],
            >>>                         output_device=args.local_rank)
        """

        def __init__(
            self,
            num_features: int,
            eps: float = 1e-5,
            momentum: float = 0.1,
            affine: bool = True,
            track_running_stats: bool = True,
            process_group: Optional[Any] = None,
            device=None,
            dtype=None
        ) -> None:
            factory_kwargs = {'device': device, 'dtype': dtype}
            super().__init__(
                num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
            )
            self.process_group = process_group

        def _check_input_dim(self, input):
            if input.dim() < 2:
                raise ValueError(
                    f"expected at least 2D input (got {input.dim()}D input)"
                )

        def _check_non_zero_input_channels(self, input):
            if input.size(1) == 0:
                raise ValueError(
                    "SyncBatchNorm number of input channels should be non-zero"
                )

        def forward(self, input: Tensor) -> Tensor:
            self._check_input_dim(input)
            self._check_non_zero_input_channels(input)

            # exponential_average_factor is set to self.momentum
            # (when it is available) only so that it gets updated
            # in ONNX graph when this node is exported to ONNX.
            if self.momentum is None:
                exponential_average_factor = 0.0
            else:
                exponential_average_factor = self.momentum

            if self.training and self.track_running_stats:
                assert self.num_batches_tracked is not None
                self.num_batches_tracked.add_(1)
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / self.num_batches_tracked.item()
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

            r"""
            Decide whether the mini-batch stats should be used for normalization rather than the buffers.
            Mini-batch stats are used in training mode, and in eval mode when buffers are None.
            """
            if self.training:
                bn_training = True
            else:
                bn_training = (self.running_mean is None) and (self.running_var is None)

            r"""
            Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
            passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
            used for normalization (i.e. in eval mode when buffers are not None).
            """
            # If buffers are not to be tracked, ensure that they won't be updated
            running_mean = (
                self.running_mean if not self.training or self.track_running_stats else None
            )
            running_var = (
                self.running_var if not self.training or self.track_running_stats else None
            )

            # Don't sync batchnorm stats in inference mode (model.eval()).
            need_sync = (bn_training and self.training and
                         torch.distributed.is_available() and torch.distributed.is_initialized())
            if need_sync:
                # currently only GPU/PrivateUse1 input is supported
                if input.device.type not in ["cuda", torch._C._get_privateuse1_backend_name()]:
                    raise ValueError("SyncBatchNorm expected input tensor to be on GPU or "
                                     f"{torch._C._get_privateuse1_backend_name()}")

                process_group = torch.distributed.group.WORLD
                if self.process_group:
                    process_group = self.process_group
                world_size = torch.distributed.get_world_size(process_group)
                need_sync = world_size > 1

            # fallback to framework BN when synchronization is not necessary
            if not need_sync:
                return F.batch_norm(
                    input,
                    running_mean,
                    running_var,
                    self.weight,
                    self.bias,
                    bn_training,
                    exponential_average_factor,
                    self.eps,
                )
            else:
                assert bn_training
                return sync_batch_norm.apply(
                    input,
                    self.weight,
                    self.bias,
                    running_mean,
                    running_var,
                    self.eps,
                    exponential_average_factor,
                    process_group,
                    world_size,
                )

        @classmethod
        def convert_sync_batchnorm(cls, module, process_group=None):
            r"""Helper function to convert all :attr:`BatchNorm*D` layers in the model to
            :class:`torch.nn.SyncBatchNorm` layers.

            Args:
                module (nn.Module): module containing one or more :attr:`BatchNorm*D` layers
                process_group (optional): process group to scope synchronization,
                    default is the whole world

            Returns:
                The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`
                layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,
                a new :class:`torch.nn.SyncBatchNorm` layer object will be returned
                instead.

            Example::

                >>> # Network with nn.BatchNorm layer
                >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
                >>> module = torch.nn.Sequential(
                >>>            torch.nn.Linear(20, 100),
                >>>            torch.nn.BatchNorm1d(100),
                >>>          ).cuda()
                >>> # creating process group (optional)
                >>> # ranks is a list of int identifying rank ids.
                >>> ranks = list(range(8))
                >>> r1, r2 = ranks[:4], ranks[4:]
                >>> # Note: every rank calls into new_group for every
                >>> # process group created, even if that rank is not
                >>> # part of the group.
                >>> # xdoctest: +SKIP("distributed")
                >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
                >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
                >>> sync_bn_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)

            """
            module_output = module
            if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
                module_output = torch.nn.SyncBatchNorm(
                    module.num_features,
                    module.eps,
                    module.momentum,
                    module.affine,
                    module.track_running_stats,
                    process_group,
                )
                if module.affine:
                    with torch.no_grad():
                        module_output.weight = module.weight
                        module_output.bias = module.bias
                module_output.running_mean = module.running_mean
                module_output.running_var = module.running_var
                module_output.num_batches_tracked = module.num_batches_tracked
                if hasattr(module, "qconfig"):
                    module_output.qconfig = module.qconfig
            for name, child in module.named_children():
                module_output.add_module(
                    name, cls.convert_sync_batchnorm(child, process_group)
                )
            del module
            return module_output

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.