mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0__cp310-cp310-win_amd64.whl

mindspore/mint/nn/layer/normalization.py

@@ -28,7 +28,11 @@ from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.nn.cell import Cell
 from mindspore.nn.layer.normalization import LayerNormExt as LayerNorm
-from mindspore.ops import group_norm
+from mindspore.communication import get_group_size
+from mindspore.communication._comm_helper import GlobalComm
+from mindspore.ops.function import batch_norm
+
+from ._functions import _SyncBatchNorm
 
 
 class _NormBase(Cell):
@@ -43,6 +47,7 @@ class _NormBase(Cell):
                  dtype=None
                  ) -> None:
         super(_NormBase, self).__init__()
+        self.set_train()
         self.shape = ops.Shape()
         self.num_features = num_features
         self.eps = eps
@@ -55,8 +60,6 @@ class _NormBase(Cell):
                 Tensor(np.empty(num_features), dtype=self.dtype), name="weight")
             self.bias = Parameter(
                 Tensor(np.empty(num_features), dtype=self.dtype), name="bias")
-            self.weight: Optional[Parameter]
-            self.bias: Optional[Parameter]
         else:
             self.weight = None
             self.bias = None
@@ -65,11 +68,8 @@ class _NormBase(Cell):
                                           requires_grad=False, name="running_mean")
             self.running_var = Parameter(Tensor(np.ones(num_features), dtype=self.dtype),
                                          requires_grad=False, name="running_var")
-            self.running_mean: Optional[Tensor]
-            self.running_var: Optional[Tensor]
-            self.num_batches_tracked = Parameter(Tensor(0, dtype=ms.float32),
+            self.num_batches_tracked = Parameter(Tensor(0, dtype=ms.int64),
                                                  requires_grad=False, name="num_batches_tracked")
-            self.num_batches_tracked: Optional[Tensor]
         else:
             self.running_mean = None
             self.running_var = None
@@ -84,7 +84,7 @@ class _NormBase(Cell):
                 np.zeros(self.num_features), dtype=self.dtype)
             one_running_var = Tensor(
                 np.ones(self.num_features), dtype=self.dtype)
-            zero_num_batches_tracked = Tensor(0, dtype=ms.float32)
+            zero_num_batches_tracked = Tensor(0, dtype=ms.int64)
 
             ops.assign(self.running_mean, zero_running_mean)
             ops.assign(self.running_var, one_running_var)
@@ -122,6 +122,7 @@ class _BatchNorm(_NormBase):
                                          dtype)
         self.training = True
 
+
     def _check_input_dim(self, input):
         raise NotImplementedError
 
@@ -135,11 +136,9 @@ class _BatchNorm(_NormBase):
 
         if self.training and self.track_running_stats:
             if self.num_batches_tracked is not None:
-                num_batches_tracked_one = Tensor(1, dtype=ms.float32)
-                ops.assign_add(self.num_batches_tracked,
-                               num_batches_tracked_one)
+                self.num_batches_tracked += 1
                 if self.momentum is None:
-                    exponential_average_factor = float(1.0 / self.num_batches_tracked)
+                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                 else:
                     exponential_average_factor = self.momentum
 
@@ -206,7 +205,7 @@ class BatchNorm1d(_BatchNorm):
         Tensor, has the same type and shape as `input`.
 
     Raises:
-        TypeError: If `num_features` is not a int number.
+        TypeError: If `num_features` is not an int number.
         TypeError: If `eps` is not a float.
         ValueError: If `num_features` is less than 1.
 
@@ -241,7 +240,7 @@ class BatchNorm2d(_BatchNorm):
 
     .. math::
 
-        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+        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
@@ -249,8 +248,8 @@ class BatchNorm2d(_BatchNorm):
     elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
 
     .. warning::
-        This API does not support Dynamic Rank.
-        This is an experimental API that is subject to change or deletion.
+        - This API does not support Dynamic Rank.
+        - This is an experimental API that is subject to change or deletion.
 
     Args:
         num_features (int): `C` from an expected input of shape :math:`(N, C, H, W)`.
@@ -263,7 +262,7 @@ class BatchNorm2d(_BatchNorm):
         track_running_stats (bool, optional): a boolean value that when set to ``True``, this
             cell tracks the running mean and variance, and when set to ``False``,
             this cell does not track such statistics. And this cell always uses batch statistics
-            in both training and eval modes. Default: ``True`` .
+            in both train and eval modes. Default: ``True`` .
         dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
 
     Inputs:
@@ -273,7 +272,7 @@ class BatchNorm2d(_BatchNorm):
         Tensor, has the same type and shape as `input`.
 
     Raises:
-        TypeError: If `num_features` is not a int number.
+        TypeError: If `num_features` is not an int number.
         TypeError: If `eps` is not a float.
         ValueError: If `num_features` is less than 1.
 
@@ -311,7 +310,7 @@ class BatchNorm3d(_BatchNorm):
 
     .. math::
 
-        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+        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
@@ -343,7 +342,7 @@ class BatchNorm3d(_BatchNorm):
         Tensor, has the same type and shape as `input`.
 
     Raises:
-        TypeError: If `num_features` is not a int number.
+        TypeError: If `num_features` is not an int number.
         TypeError: If `eps` is not a float.
         ValueError: If `num_features` is less than 1.
 
@@ -402,7 +401,7 @@ class GroupNorm(Cell):
           additional dimensions.
 
     Outputs:
-        Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `x`.
+        Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `input`.
 
     Raises:
         TypeError: If `num_groups` or `num_channels` is not an int.
@@ -435,8 +434,8 @@ class GroupNorm(Cell):
         """Initialize GroupNorm."""
         super(GroupNorm, self).__init__()
         ms_dtype = mstype.float32 if dtype is None else dtype
-        gamma_init = 'ones'
-        beta_init = 'zeros'
+        weight_init = 'ones'
+        bias_init = 'zeros'
 
         self.num_groups = validator.check_positive_int(
             num_groups, "num_groups", self.cls_name)
@@ -450,14 +449,14 @@ class GroupNorm(Cell):
         self.affine = validator.check_bool(
             affine, arg_name="affine", prim_name=self.cls_name)
 
-        self.gamma = Parameter(initializer(
-            gamma_init, self.num_channels, dtype=ms_dtype), name="gamma", requires_grad=affine)
-        self.beta = Parameter(initializer(
-            beta_init, self.num_channels, dtype=ms_dtype), name="beta", requires_grad=affine)
+        self.weight = Parameter(initializer(
+            weight_init, self.num_channels, dtype=ms_dtype), name="weight", requires_grad=affine)
+        self.bias = Parameter(initializer(
+            bias_init, self.num_channels, dtype=ms_dtype), name="bias", requires_grad=affine)
 
     def _cal_output(self, x):
         """calculate groupnorm output"""
-        return group_norm(x, self.num_groups, self.gamma, self.beta, self.eps)
+        return ops.group_norm(x, self.num_groups, self.weight, self.bias, self.eps)
 
     def extend_repr(self):
         return 'num_groups={}, num_channels={}, eps={}, affine={}'.format(
@@ -468,10 +467,206 @@ class GroupNorm(Cell):
         return output
 
 
+class SyncBatchNorm(_BatchNorm):
+    r"""
+    Sync Batch Normalization layer over a N-dimension input.
+
+    Sync Batch Normalization is cross device synchronized Batch Normalization. The implementation of Batch
+    Normalization only normalizes the data within each device. Sync Batch Normalization will normalize the input
+    within the group. It has been described in the paper `Batch Normalization: Accelerating Deep Network Training by
+    Reducing Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`_. It rescales and recenters the
+    feature using a mini-batch of data and the learned parameters which can be described in the following formula.
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    .. warning::
+            This is an experimental API that is subject to change or deletion.
+
+    Args:
+        num_features (int): `C` from an expected input of size :math:`(N, C, +)`.
+        eps (float, optional): :math:`\epsilon`, a value added to the denominator for numerical stability.
+            Default: ``1e-5`` .
+        momentum (float, optional): A floating hyperparameter of the momentum for the
+            running_mean and running_var computation. Default: ``0.1`` .
+        affine (bool, optional): A bool value. When set to ``True`` , :math:`\gamma` and :math:`\beta` are learnable
+            parameters. When set to ``False`` , :math:`\gamma` and :math:`\beta` are unlearnable parameters.
+            Default: ``True`` .
+        track_running_stats (bool, optional): a boolean value that when set to ``True``, this
+            cell tracks the running mean and variance, and when set to ``False``,
+            this cell does not track such statistics. And this cell always uses batch statistics
+            in both training and eval modes. Default: ``True`` .
+        process_group (str, optional): synchronization of stats happen within each process group individually.
+            Default behavior is synchronization across the whole world. Default: ``None`` .
+        dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **x** (Tensor) - Tensor of shape :math:`(N, C_{in}, +)`.
+
+    Outputs:
+        Tensor, the normalized, scaled, offset tensor, of shape :math:`(N, C_{out}, +)`.
+
+    Raises:
+        TypeError: If `num_features` is not an int.
+        TypeError: If `eps` is not a float.
+        ValueError: If `num_features` is less than 1.
+        ValueError: If `momentum` is not in range [0, 1].
+        ValueError: If rank_id in `process_group` is not in range [0, rank_size).
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
+            Here, examples use msrun to pull multi-process distributed tasks across nodes with a single command
+            line instruction.
+            Please see the `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with multiple devices.
+
+        >>> # Firstly, preparing test_syncbn.py:
+        >>> import numpy as np
+        >>> import mindspore
+        >>> import mindspore.context as context
+        >>> from mindspore.mint.nn.layer import SyncBatchNorm
+        >>> from mindspore import Tensor
+        >>> from mindspore.communication import init, create_group, get_local_rank
+        >>> init()
+        >>> context.set_context(mode=context.PYNATIVE_MODE, device_target="Ascend")
+        >>> group = "0-1"
+        >>> rank_ids = [0, 1]
+        >>> create_group(group, rank_ids)
+        >>> sync_batch_norm = SyncBatchNorm(num_features=2, process_group=group, dtype=mindspore.float32)
+        >>> sync_batch_norm.set_train(False)
+        >>> input_x = Tensor(np.linspace(0, 5, 2*2*2*2), mindspore.float32).reshape(2, 2, 2, 2)
+        >>> output_data = sync_batch_norm(input_x)
+        >>> # Then, executing the command such as the following:
+        >>> # msrun --worker_num=2 --local_worker_num=2 --master_port=8975 --log_dir=msrun_log --join=True
+        >>> # --cluster_time_out=100 pytest -s -v test_syncbn.py
+
+    """
+    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[str] = None,
+                 dtype=None):
+        super(SyncBatchNorm, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats, dtype
+        )
+
+        self.process_group = process_group if process_group else GlobalComm.WORLD_COMM_GROUP
+        self.world_size = get_group_size(self.process_group)
+        self.sync_batch_norm = _SyncBatchNorm(
+            self.num_features, self.world_size, self.dtype)
+
+    def _check_input_dim(self, input):
+        if input.ndim < 2:
+            raise ValueError(
+                "expected at least 2D input (got {}D input)".format(input.ndim)
+            )
+
+    def _check_non_zero_input_channels(self, input):
+        if input.shape[1] == 0:
+            raise ValueError(
+                "SyncBatchNorm number of input channels should be non-zero"
+            )
+
+    def construct(self, input: Tensor) -> Tensor:
+        # currently only GPU input is supported
+
+        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:
+            self.num_batches_tracked += 1
+            if self.momentum is None:  # use cumulative moving average
+                exponential_average_factor = 1.0 / float(self.num_batches_tracked.value())
+            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)
+        if need_sync:
+            need_sync = self.world_size > 1
+
+        # fallback to framework BN when synchronization is not necessary
+        if not need_sync:
+            if self.weight is None:
+                weight = Tensor(np.ones(self.num_features), dtype=self.dtype)
+            else:
+                weight = self.weight
+            if self.bias is None:
+                bias = Tensor(np.zeros(self.num_features), dtype=self.dtype)
+            else:
+                bias = self.bias
+            if running_mean is None or running_var is None:
+                raise ValueError(
+                    "running mean or running var can\'t be none for batch_norm.")
+            return batch_norm(input,
+                              running_mean,
+                              running_var,
+                              weight,
+                              bias,
+                              bn_training,
+                              exponential_average_factor,
+                              self.eps)
+        else:
+            output = self.sync_batch_norm(input,
+                                          self.weight,
+                                          self.bias,
+                                          running_mean,
+                                          running_var,
+                                          self.eps,
+                                          exponential_average_factor,
+                                          self.process_group,
+                                          self.world_size)
+            return output
+
+
 __all__ = [
     'GroupNorm',
     'BatchNorm1d',
     'BatchNorm2d',
     'BatchNorm3d',
     'LayerNorm',
+    'SyncBatchNorm',
 ]