mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/nn/layer/normalization.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2021 Huawei Technologies Co., Ltd
+# Copyright 2020-2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@ from __future__ import division
 
 import itertools
 import numbers
+import hashlib
 
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
@@ -25,21 +26,26 @@ from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common.parameter import Parameter
 from mindspore.common.initializer import initializer, Initializer
 from mindspore.common.tensor import Tensor
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import constexpr, _primexpr
 import mindspore.context as context
-from mindspore._checkparam import Rel
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore._extends import cell_attr_register
 from mindspore.communication.management import get_group_size, get_rank
 from mindspore.communication import management
 from mindspore.common import dtype as mstype
 from mindspore.parallel._utils import _is_in_auto_parallel_mode
 from mindspore.nn.cell import Cell
+from mindspore import log as logger
 
 __all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LayerNorm', 'GroupNorm',
            'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d']
 
-SYNC_BN_GROUP_NAME = ""
+
+def _check_dim(val, target, cls_name):
+    def _check(val, target, cls_name):
+        if val != target:
+            raise ValueError(f"For '{cls_name}', the in_shape must have {target} dims, but got {val}.")
+    _check(val, target, cls_name)
 
 
 class _BatchNorm(Cell):
@@ -121,11 +127,13 @@ class _BatchNorm(Cell):
         self.assign_sub_mean = P.AssignSub().shard(data_parallel_strategy)
         self.assign_sub_var = P.AssignSub().shard(data_parallel_strategy)
 
+
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         raise NotImplementedError
 
+
     def construct(self, x):
         self._check_input_dim(self.shape(x), self.cls_name)
         if self.use_batch_statistics is None:
@@ -164,7 +172,7 @@ class _BatchNorm(Cell):
 class BatchNorm1d(_BatchNorm):
     r"""
     This layer
-    applies Batch Normalization over a 2D input (a mini-batch of 1D inputs) to
+    applies Batch Normalization over a 2D or 3D input (a mini-batch of 1D or 2D inputs) to
     reduce internal covariate shift. Batch Normalization is widely used in convolutional networks.
     For the setailed contents, refer to `Batch Normalization: Accelerating Deep Network Training by
     Reducing Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`_. It
@@ -179,14 +187,14 @@ class BatchNorm1d(_BatchNorm):
         recommended to be changed after net was initialized.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C).
-        eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
+        num_features (int): number of features or channels `C` of the input `x` .
+        eps (float): :math:`\epsilon` added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.9.
-        affine (bool): A bool value. When set to True, gamma and beta can be learned. Default: True.
-        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the gamma weight.
+        affine (bool): A bool value. When set to True, :math:`\gamma` and :math:`\beta` can be learned. Default: True.
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'ones'.
-        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'zeros'.
         moving_mean_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving mean.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'zeros'.
@@ -200,10 +208,11 @@ class BatchNorm1d(_BatchNorm):
             Default: 'NCHW'.
 
     Inputs:
-        - **x** (Tensor) - Tensor of shape :math:`(N, C_{in})`.
+        - **x** (Tensor) - Tensor of shape :math:`(N, C)` or :math:`(N, C, L)` ,
+          where `N` is the batch size, `C` is the number of features or channels, and `L` is the sequence length.
 
     Outputs:
-        Tensor, the normalized, scaled, offset tensor, of shape :math:`(N, C_{out})`.
+        Tensor, the normalized, scaled, offset tensor, of shape :math:`(N, C)` or :math:`(N, C, L)` .
 
     Raises:
         TypeError: If `num_features` is not an int.
@@ -228,11 +237,13 @@ class BatchNorm1d(_BatchNorm):
     """
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
+        def _check(dim):
+            if dim not in (2, 3):
+                raise ValueError(f"For '{cls_name}', the must have 2 dims or 3 dims, but got {dim}.")
         dim = len(shape)
-        if dim != 2:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 2 dims, but got {dim}.")
+        _check(dim)
 
 
 class BatchNorm2d(_BatchNorm):
@@ -254,22 +265,22 @@ class BatchNorm2d(_BatchNorm):
         Note that the formula for updating the :math:`moving\_mean` and :math:`moving\_var` is
 
         .. math::
-            \text{moving_mean}=\text{moving_mean∗momentum}+μ_β\text{∗(1−momentum)}\\
-            \text{moving_var}=\text{moving_var∗momentum}+σ^2_β\text{∗(1−momentum)}
+            \text{moving_mean}=\text{moving_mean*momentum}+μ_β\text{*(1−momentum)}\\
+            \text{moving_var}=\text{moving_var*momentum}+σ^2_β\text{*(1−momentum)}
 
         where :math:`moving\_mean` is the updated mean, :math:`moving\_var` is the updated variance,
         :math:`μ_β, σ^2_β` are the observed value (mean and variance) of each batch of data.
 
     Args:
-        num_features (int): The number of channels of the input tensor. Expected input size is (N, C, H, W),
+        num_features (int): The number of channels of the input tensor. Expected input size is :math:`(N, C, H, W)`,
             `C` represents the number of channels.
-        eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
+        eps (float): :math:`\epsilon` added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.9.
-        affine (bool): A bool value. When set to True, gamma and beta can be learned. Default: True.
-        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the gamma weight.
+        affine (bool): A bool value. When set to True, :math:`\gamma` and :math:`\beta` can be learned. Default: True.
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'ones'.
-        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'zeros'.
         moving_mean_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving mean.
             The values of str refer to the function `initializer` including 'zeros', 'ones', etc. Default: 'zeros'.
@@ -288,10 +299,10 @@ class BatchNorm2d(_BatchNorm):
             Default: 'NCHW'.
 
     Inputs:
-        - **x** (Tensor) - Tensor of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
+        - **x** (Tensor) - Tensor of shape :math:`(N, C, H, W)`.
 
     Outputs:
-        Tensor, the normalized, scaled, offset tensor, of shape :math:`(N, C_{out}, H_{out}, W_{out})`.
+        Tensor, the normalized, scaled, offset tensor, of shape :math:`(N, C, H, W)`.
 
     Raises:
         TypeError: If `num_features` is not an int.
@@ -320,11 +331,10 @@ class BatchNorm2d(_BatchNorm):
     """
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 4:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 4 dims, but got {dim}.")
+        _check_dim(dim, 4, cls_name)
 
 
 class BatchNorm3d(Cell):
@@ -344,7 +354,7 @@ class BatchNorm3d(Cell):
         where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C, D, H, W).
+        num_features (int): `C` from an expected input of size :math:`(N, C, D, H, W)` .
         eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.9.
@@ -414,11 +424,11 @@ class BatchNorm3d(Cell):
         self.reshape = P.Reshape()
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 5:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 5 dims, but got {dim}.")
+        _check_dim(dim, 5, cls_name)
+
 
     def construct(self, x):
         x_shape = self.shape(x)
@@ -429,6 +439,16 @@ class BatchNorm3d(Cell):
         return bn3d_out
 
 
+SYNCBN_GROUP_DICT = None
+
+
+def _syncbatchnorm_group_dict():
+    global SYNCBN_GROUP_DICT
+    if SYNCBN_GROUP_DICT is None:
+        SYNCBN_GROUP_DICT = dict()
+    return SYNCBN_GROUP_DICT
+
+
 class SyncBatchNorm(_BatchNorm):
     r"""
     Sync Batch Normalization layer over a N-dimension input.
@@ -446,15 +466,16 @@ class SyncBatchNorm(_BatchNorm):
         Currently, SyncBatchNorm only supports 2D and 4D inputs.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C, H, W).
-        eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
+        num_features (int): `C` from an expected input of size :math:`(N, C, H, W)`.
+        eps (float): :math:`\epsilon`, a value added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.9.
-        affine (bool): A bool value. When set to True, gamma and beta can be learned. Default: True.
-        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the gamma weight.
+        affine (bool): A bool value. When set to True, :math:`\gamma` and :math:`\beta` can be learned.
+            Default: True.
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
             'he_uniform', etc. Default: 'ones'.
-        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
             'he_uniform', etc. Default: 'zeros'.
         moving_mean_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving mean.
@@ -495,11 +516,11 @@ class SyncBatchNorm(_BatchNorm):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with multiple devices.
 
@@ -525,7 +546,7 @@ class SyncBatchNorm(_BatchNorm):
           [[ 0.999995 0.999995 ]
            [ 0.999995 0.999995 ]]]]
     """
-
+    @cell_attr_register(attrs=['num_features', 'process_groups'])
     def __init__(self,
                  num_features,
                  eps=1e-5,
@@ -548,7 +569,7 @@ class SyncBatchNorm(_BatchNorm):
                                             moving_var_init,
                                             use_batch_statistics)
         self.is_global = False
-        global SYNC_BN_GROUP_NAME
+        self.group_name = None
         self.process_groups = process_groups
         if self.process_groups != 0:
             self.rank_id = get_rank()
@@ -560,43 +581,53 @@ class SyncBatchNorm(_BatchNorm):
             elif self.rank_size > 1:
                 self.is_global = True
                 self.group_device_num = self.rank_size
-                self.device_list = [i for i in range(0, self.rank_size)]
                 if context.get_context("device_target") == "Ascend":
-                    if SYNC_BN_GROUP_NAME == "":
-                        SYNC_BN_GROUP_NAME = "sync_bn_group0"
-                        management.create_group(SYNC_BN_GROUP_NAME, self.device_list)
+                    self.group_name = "hccl_world_group"
                 elif context.get_context("device_target") == "GPU":
-                    if SYNC_BN_GROUP_NAME == "":
-                        SYNC_BN_GROUP_NAME = "nccl_world_group"
+                    self.group_name = "nccl_world_group"
 
         if self.is_global:
             self.bn_train = inner.SyncBatchNorm(epsilon=self.eps,
                                                 momentum=self.momentum,
-                                                group=SYNC_BN_GROUP_NAME,
+                                                group=self.group_name,
                                                 device_num=self.group_device_num)
 
     def _create_sync_groups(self):
-        for i in range(len(self.process_groups)):
-            validator.check_isinstance("process_groups[%d]" % i, self.process_groups[i], list)
-            self.group_device_num = len(self.process_groups[i])
-            if self.rank_id in self.process_groups[i] and self.group_device_num > 1:
+        """ create groups by process groups. """
+        for sub_group in self.process_groups:
+            validator.check_isinstance("sub group", sub_group, list)
+            self.group_device_num = len(sub_group)
+            if self.rank_id in sub_group and self.group_device_num > 1:
                 self.is_global = True
-                global SYNC_BN_GROUP_NAME
-                if SYNC_BN_GROUP_NAME == "":
-                    SYNC_BN_GROUP_NAME = "sync_bn_group%d" % i
-                    management.create_group(SYNC_BN_GROUP_NAME, self.process_groups[i])
+                rank_list_name = '_'.join('%s' % id for id in sub_group)
+                group_dict = _syncbatchnorm_group_dict()
+                if rank_list_name not in group_dict:
+                    md5 = hashlib.md5()
+                    md5.update(rank_list_name.encode('utf-8'))
+                    hash_name = md5.hexdigest()
+                    self.group_name = str(self.group_device_num) + '_' + hash_name
+                    group_dict[rank_list_name] = self.group_name
+                    management.create_group(self.group_name, sub_group)
+                    logger.info("create group for sync batchnorm, the rank list is {}, the group name is {}".format(
+                        rank_list_name, self.group_name))
+                else:
+                    self.group_name = group_dict[rank_list_name]
+                    logger.info("the group for {} already exists, no need to create".format(rank_list_name))
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
+        def _check(dim):
+            if dim not in (2, 4):
+                raise ValueError(f"For '{cls_name}', the must have 2 dims or 4 dims, but got {dim}.")
         dim = len(shape)
-        if dim not in (2, 4):
-            raise ValueError(f"For '{cls_name}', the must have 2 dims or 4 dims, but got {dim}.")
+        _check(dim)
+
 
     def _check_rank_ids(self, process_groups, rank_size):
         seen = set()
         for rid in itertools.chain(*process_groups):
-            validator.check_int_range(rid, 0, rank_size, Rel.INC_LEFT, "rank id in process_groups", self.cls_name)
+            validator.check_int_range(rid, 0, rank_size, validator.INC_LEFT, "rank id in process_groups", self.cls_name)
             if rid in seen:
                 raise ValueError(f"For '{self.cls_name}', rank id in 'process_groups' must not be duplicated, "
                                  f"but got {process_groups}.")
@@ -625,13 +656,13 @@ class LayerNorm(Cell):
         begin_params_axis (int): The first parameter(beta, gamma)dimension: scale and centering parameters
             will have dimensions `begin_params_axis: rank(inputs)` and will be broadcast with
             the normalized inputs accordingly, the value should be in [-1, rank(input)). Default: -1.
-        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the gamma weight.
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
             'he_uniform', etc. Default: 'ones'.
-        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
             'he_uniform', etc. Default: 'zeros'.
-        epsilon (float): A value added to the denominator for numerical stability. Default: 1e-7.
+        epsilon (float): :math:`\epsilon` added to the denominator for numerical stability. Default: 1e-7.
 
     Inputs:
         - **x** (Tensor) - The shape of `x` is :math:`(x_1, x_2, ..., x_R)`,
@@ -775,7 +806,7 @@ class InstanceNorm1d(_InstanceNorm):
         where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C, L).
+        num_features (int): `C` from an expected input of size :math:`(N, C, L)`.
         eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.1.
@@ -823,11 +854,11 @@ class InstanceNorm1d(_InstanceNorm):
     """
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 3:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 3 dims, but got {dim}.")
+        _check_dim(dim, 3, cls_name)
+
 
 
 class InstanceNorm2d(_InstanceNorm):
@@ -854,7 +885,7 @@ class InstanceNorm2d(_InstanceNorm):
         where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C, H, W).
+        num_features (int): `C` from an expected input of size :math:`(N, C, H, W)`.
         eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.1.
@@ -902,11 +933,10 @@ class InstanceNorm2d(_InstanceNorm):
     """
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 4:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 4 dims, but got {dim}.")
+        _check_dim(dim, 4, cls_name)
 
 
 class InstanceNorm3d(_InstanceNorm):
@@ -933,7 +963,7 @@ class InstanceNorm3d(_InstanceNorm):
         where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value.
 
     Args:
-        num_features (int): `C` from an expected input of size (N, C, D, H, W).
+        num_features (int): `C` from an expected input of size :math:`(N, C, D, H, W)`.
         eps (float): A value added to the denominator for numerical stability. Default: 1e-5.
         momentum (float): A floating hyperparameter of the momentum for the
             running_mean and running_var computation. Default: 0.1.
@@ -979,12 +1009,12 @@ class InstanceNorm3d(_InstanceNorm):
         >>> print(output.shape)
         (2, 3, 5, 2, 2)
     """
+
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 5:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 5 dims, but got {dim}.")
+        _check_dim(dim, 5, cls_name)
 
 
 class GroupNorm(Cell):
@@ -1007,10 +1037,10 @@ class GroupNorm(Cell):
         affine (bool): A bool value, this layer will have learnable affine parameters when set to true. Default: True.
         gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the gamma weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
-            'he_uniform', etc. Default: 'ones'. If gamma_init is a Tensor, the shape must be [num_channels].
+            'he_uniform', etc. Default: 'ones'. If gamma_init is a Tensor, the shape must be :math:`(num\_channels)`.
         beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
             The values of str refer to the function `initializer` including 'zeros', 'ones', 'xavier_uniform',
-            'he_uniform', etc. Default: 'zeros'. If beta_init is a Tensor, the shape must be [num_channels].
+            'he_uniform', etc. Default: 'zeros'. If beta_init is a Tensor, the shape must be :math:`(num\_channels)`.
 
     Inputs:
         - **x** (Tensor) - The input feature with shape :math:`(N, C, H, W)` .
@@ -1079,19 +1109,20 @@ class GroupNorm(Cell):
         return output
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _check_input_dim(shape, cls_name):
         dim = len(shape)
-        if dim != 4:
-            raise ValueError(f"For '{cls_name}', the in_shape must have 4 dims, but got {dim}.")
+        _check_dim(dim, 4, cls_name)
 
     @staticmethod
-    @constexpr
+    @_primexpr
     def _channel_check(channel, num_channel, prim_name=None):
-        msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-        if channel != num_channel:
-            raise ValueError(f"{msg_prefix} channel(the second dim of the input 'x') must be equal to num_channels, "
-                             f"but got channel: {channel}, num_channels: {num_channel}.")
+        def _check():
+            msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
+            if channel != num_channel:
+                raise ValueError(f"{msg_prefix} channel(the second dim of the input 'x') must be equal to "
+                                 f"num_channels, but got channel: {channel}, num_channels: {num_channel}.")
+        _check()
 
     @staticmethod
     @constexpr