mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-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,22 +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.common._decorator import deprecated
-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',
-           'GlobalBatchNorm', 'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d']
+           '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):
@@ -57,9 +62,6 @@ class _BatchNorm(Cell):
                  moving_mean_init='zeros',
                  moving_var_init='ones',
                  use_batch_statistics=None,
-                 device_num_each_group=1,
-                 process_groups=0,
-                 input_dims='2d',
                  data_format='NCHW'):
         """Initialize _BatchNorm."""
         super(_BatchNorm, self).__init__()
@@ -70,7 +72,6 @@ class _BatchNorm(Cell):
         if momentum < 0 or momentum > 1:
             raise ValueError(f"For '{self.cls_name}', the 'momentum' must be a number in range [0, 1], "
                              f"but got {momentum}.")
-        self.input_dims = input_dims
         self.format = validator.check_string(data_format, ['NCHW', 'NHWC'], 'format', self.cls_name)
         if context.get_context("device_target") != "GPU" and self.format == "NHWC":
             raise ValueError(f"For '{self.cls_name}', the 'NHWC' format only support in GPU target, but got device "
@@ -93,39 +94,8 @@ class _BatchNorm(Cell):
             gamma_init, num_features), name="gamma", requires_grad=affine)
         self.beta = Parameter(initializer(
             beta_init, num_features), name="beta", requires_grad=affine)
-        self.group_device_num = validator.check_positive_int(device_num_each_group, "device_num_each_group",
-                                                             self.cls_name)
-        self.process_groups = process_groups
-        self.is_global = False
+
         self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
-        global SYNC_BN_GROUP_NAME
-        # for GlobalBatchNorm
-        if self.group_device_num != 1:
-            self.rank_id = get_rank()
-            self.rank_size = get_group_size()
-            self.device_list = [i for i in range(0, self.rank_size)]
-            self.rank_list = self.list_group(self.device_list, self.group_device_num)
-            self.rank_list_idx = len(self.rank_list)
-            self._create_global_groups()
-        # for SyncBatchNorm
-        if self.process_groups != 0:
-            self.rank_id = get_rank()
-            self.rank_size = get_group_size()
-            if self.process_groups is not None:
-                validator.check_isinstance("process_groups", self.process_groups, list)
-                self._check_rank_ids(self.process_groups, self.rank_size)
-                self._create_sync_groups()
-            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)
-                elif context.get_context("device_target") == "GPU":
-                    if SYNC_BN_GROUP_NAME == "":
-                        SYNC_BN_GROUP_NAME = "nccl_world_group"
 
         self.shape = P.Shape()
         self.reduce_mean = P.ReduceMean(keep_dims=True)
@@ -137,20 +107,11 @@ class _BatchNorm(Cell):
         self._target = context.get_context("device_target")
         self.is_graph_mode = context.get_context("mode") == context.GRAPH_MODE
         self.momentum = 1.0 - momentum
-        if context.get_context("enable_ge"):
-            self.is_ge_backend = True
-        else:
-            self.is_ge_backend = False
 
         self.bn_train = P.BatchNorm(is_training=True,
                                     epsilon=self.eps,
                                     momentum=self.momentum,
                                     data_format=self.format)
-        if self.is_global:
-            self.bn_train = inner.SyncBatchNorm(epsilon=self.eps,
-                                                momentum=self.momentum,
-                                                group=SYNC_BN_GROUP_NAME,
-                                                device_num=self.group_device_num)
 
         self.bn_infer = P.BatchNorm(is_training=False, epsilon=self.eps, data_format=self.format)
         if _is_in_auto_parallel_mode():
@@ -166,22 +127,15 @@ class _BatchNorm(Cell):
         self.assign_sub_mean = P.AssignSub().shard(data_parallel_strategy)
         self.assign_sub_var = P.AssignSub().shard(data_parallel_strategy)
 
-    def list_group(self, world_rank, group_size):
-        """ Check whether world_rank and group_size  are valid. """
-        if group_size > get_group_size():
-            raise ValueError(f"For '{self.cls_name}', the 'device_num_each_group' cannot be greater than "
-                             f"local rank size, but got 'device_num_each_group': {group_size}, "
-                             f"local rank size: {get_group_size()}.")
-        if len(world_rank) % group_size != 0:
-            raise ValueError(f"For '{self.cls_name}', the dimension of device_list must be divisible by "
-                             f"'device_num_each_group', but got the length of device_list: {len(world_rank)}, "
-                             f"'device_num_each_group': {group_size}.")
-        world_rank_list = zip(*(iter(world_rank),) * group_size)
-        group_list = [list(i) for i in world_rank_list]
-        return group_list
+
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        raise NotImplementedError
+
 
     def construct(self, x):
-        _shape_check_bn(self.shape(x), self.input_dims, self.cls_name)
+        self._check_input_dim(self.shape(x), self.cls_name)
         if self.use_batch_statistics is None:
             if self.training:
                 return self.bn_train(x,
@@ -214,98 +168,11 @@ class _BatchNorm(Cell):
             self.num_features, self.eps, 1.0 - self.momentum, self.gamma, self.beta, \
             self.moving_mean, self.moving_variance)
 
-    def _check_data_dim(self, x):
-        raise NotImplementedError
-
-    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)
-            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}.")
-            seen.add(rid)
-
-    def _create_global_groups(self):
-        for i in range(self.rank_list_idx):
-            if self.rank_id in self.rank_list[i]:
-                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.rank_list[i])
-
-    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:
-                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])
-
-
-@constexpr
-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}.")
-
-
-@constexpr
-def _shape_check(in_shape, prim_name=None):
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    if len(in_shape) != 4:
-        raise ValueError(f"{msg_prefix} in_shape must has 4 dims, but got the length of in_shape: {len(in_shape)}.")
-
-
-@constexpr
-def _shape_check_bn(in_shape, in_dims, prim_name=None):
-    """check input dims of batch norm."""
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    dim = len(in_shape)
-    if in_dims == '1d' and dim != 2:
-        raise ValueError(f"{msg_prefix} in_shape must have 2 dims, but got {len(in_shape)}.")
-    if in_dims == '2d' and dim != 4:
-        raise ValueError(f"{msg_prefix} in_shape must have 4 dims, but got {len(in_shape)}.")
-    if in_dims == '3d' and dim != 5:
-        raise ValueError(f"{msg_prefix} in_shape must have 5 dims, but got {len(in_shape)}.")
-    if in_dims == 'both' and dim != 2 and dim != 4:
-        raise ValueError(f"{msg_prefix} in_shape must have 2 dims or 4 dims, but got {len(in_shape)}.")
-
-
-@constexpr
-def _shape_check_in(in_shape, in_dims, prim_name=None):
-    """check input dims of batch norm."""
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    dim = len(in_shape)
-    if in_dims == '1d' and dim != 3:
-        raise ValueError(f"{msg_prefix} in_shape must have 3 dims, but got {len(in_shape)}.")
-    if in_dims == '2d' and dim != 4:
-        raise ValueError(f"{msg_prefix} in_shape must have 4 dims, but got {len(in_shape)}.")
-    if in_dims == '3d' and dim != 5:
-        raise ValueError(f"{msg_prefix} in_shape must have 5 dims, but got {len(in_shape)}.")
-
-
-@constexpr
-def _shape_infer(x_shape, num_feature):
-    """global Batch Normalization shape and axes infer"""
-    if len(x_shape) == 4:
-        axes = (0, 2, 3)
-        re_shape = (1, num_feature, 1, 1)
-    else:
-        axes = (0,)
-        re_shape = (1, num_feature)
-    return axes, re_shape
-
 
 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
@@ -320,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'.
@@ -337,12 +204,15 @@ class BatchNorm1d(_BatchNorm):
             use the mean value and variance value of specified value. If None, the training process will use the mean
             and variance of current batch data and track the running mean and variance, the evaluation process will use
             the running mean and variance. Default: None.
+        data_format (str): The optional value for data format, is 'NHWC' or 'NCHW'.
+            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.
@@ -366,31 +236,14 @@ class BatchNorm1d(_BatchNorm):
          [ 0.4999975   0.399998   0.59999704 0.89999545 ]]
     """
 
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.9,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros',
-                 moving_mean_init='zeros',
-                 moving_var_init='ones',
-                 use_batch_statistics=None):
-        """Initialize BatchNorm1d."""
-        super(BatchNorm1d, self).__init__(num_features,
-                                          eps,
-                                          momentum,
-                                          affine,
-                                          gamma_init,
-                                          beta_init,
-                                          moving_mean_init,
-                                          moving_var_init,
-                                          use_batch_statistics,
-                                          input_dims='1d')
-
-    def _check_data_dim(self, x):
-        if x.ndim != 2:
-            pass
+    @staticmethod
+    @_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)
+        _check(dim)
 
 
 class BatchNorm2d(_BatchNorm):
@@ -412,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'.
@@ -446,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.
@@ -477,46 +330,11 @@ class BatchNorm2d(_BatchNorm):
            [ 0.999995 0.999995 ]]]]
     """
 
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.9,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros',
-                 moving_mean_init='zeros',
-                 moving_var_init='ones',
-                 use_batch_statistics=None,
-                 data_format='NCHW'):
-        """Initialize BatchNorm2d."""
-        super(BatchNorm2d, self).__init__(num_features,
-                                          eps,
-                                          momentum,
-                                          affine,
-                                          gamma_init,
-                                          beta_init,
-                                          moving_mean_init,
-                                          moving_var_init,
-                                          use_batch_statistics,
-                                          input_dims='2d',
-                                          data_format=data_format)
-
-    def _check_data_dim(self, x):
-        if x.ndim != 4:
-            pass
-
-
-@constexpr
-def _check_3d_shape(input_shape, prim_name=None):
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    if len(input_shape) != 5:
-        raise ValueError(f"{msg_prefix} input_shape must be 5-dimensional, but got the length of input_shape: "
-                         f"{len(input_shape)}.")
-
-
-@constexpr
-def _check_dtype(dtype, valid_dtypes, args_name, prim_name=None):
-    validator.check_type_name(args_name, dtype, valid_dtypes, prim_name)
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 4, cls_name)
 
 
 class BatchNorm3d(Cell):
@@ -536,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.
@@ -553,7 +371,6 @@ class BatchNorm3d(Cell):
             use the mean value and variance value of specified value. If None, the training process will use the mean
             and variance of current batch data and track the running mean and variance, the evaluation process will use
             the running mean and variance. Default: None.
-        data_format (str): The optional value for data format is 'NCDHW'. Default: 'NCDHW'.
 
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(N, C_{in}, D_{in}, H_{in}, W_{in})`.
@@ -566,7 +383,6 @@ class BatchNorm3d(Cell):
         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 `data_format` is not 'NCDHW'.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -591,12 +407,9 @@ class BatchNorm3d(Cell):
                  beta_init='zeros',
                  moving_mean_init='zeros',
                  moving_var_init='ones',
-                 use_batch_statistics=None,
-                 data_format='NCDHW'):
+                 use_batch_statistics=None):
         """Initialize BatchNorm3d."""
         super(BatchNorm3d, self).__init__()
-        self.format = validator.check_string(data_format, ['NCDHW'], 'format', self.cls_name)
-        self.reshape = P.Reshape()
         self.bn2d = BatchNorm2d(num_features=num_features,
                                 eps=eps,
                                 momentum=momentum,
@@ -607,57 +420,33 @@ class BatchNorm3d(Cell):
                                 moving_var_init=moving_var_init,
                                 use_batch_statistics=use_batch_statistics,
                                 data_format="NCHW")
+        self.shape = P.Shape()
+        self.reshape = P.Reshape()
 
-    def construct(self, input_x):
-        x_shape = F.shape(input_x)
-        _check_3d_shape(x_shape, self.cls_name)
-        input_x = self.reshape(input_x, (x_shape[0], x_shape[1], x_shape[2] * x_shape[3], x_shape[4]))
-        bn2d_out = self.bn2d(input_x)
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 5, cls_name)
+
+
+    def construct(self, x):
+        x_shape = self.shape(x)
+        self._check_input_dim(x_shape, self.cls_name)
+        x = self.reshape(x, (x_shape[0], x_shape[1], x_shape[2] * x_shape[3], x_shape[4]))
+        bn2d_out = self.bn2d(x)
         bn3d_out = self.reshape(bn2d_out, x_shape)
         return bn3d_out
 
 
-class GlobalBatchNorm(_BatchNorm):
-    r"""
-    The GlobalBatchNorm interface is deprecated, please use the :class:`mindspore.nn.SyncBatchNorm` instead.
+SYNCBN_GROUP_DICT = None
 
-    Supported Platforms:
-        deprecated
-    """
 
-    @deprecated("1.2", "SyncBatchNorm", True)
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.9,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros',
-                 moving_mean_init='zeros',
-                 moving_var_init='ones',
-                 use_batch_statistics=None,
-                 device_num_each_group=2):
-        """Initialize GlobalBatchNorm."""
-        super(GlobalBatchNorm, self).__init__(num_features,
-                                              eps,
-                                              momentum,
-                                              affine,
-                                              gamma_init,
-                                              beta_init,
-                                              moving_mean_init,
-                                              moving_var_init,
-                                              use_batch_statistics,
-                                              device_num_each_group,
-                                              input_dims='both')
-        self.group_device_num = validator.check_positive_int(device_num_each_group, "device_num_each_group",
-                                                             self.cls_name)
-        if self.group_device_num <= 1:
-            raise ValueError(f"For '{self.cls_name}', the 'device_num_each_group' must be greater than 1, "
-                             f"but got {self.group_device_num}.")
-
-    def _check_data_dim(self, x):
-        if x.dim == 0:
-            pass
+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):
@@ -677,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.
@@ -721,10 +511,19 @@ class SyncBatchNorm(_BatchNorm):
         ``Ascend``
 
     Examples:
-        >>> # This example should be run with multiple processes.
-        >>> # Please refer to the Programming Guide > Distributed Training -> Distributed Parallel Usage Example
-        >>> # on mindspore.cn and focus on the contents of these three parts: Configuring Distributed Environment
-        >>> # Variables, Calling the Collective Communication Library, Running the Script.
+        .. 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.
+            Please see the `Ascend tutorial
+            <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/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with multiple devices.
+
         >>> import numpy as np
         >>> import mindspore as ms
         >>> from mindspore.communication import init
@@ -747,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,
@@ -768,13 +567,71 @@ class SyncBatchNorm(_BatchNorm):
                                             beta_init,
                                             moving_mean_init,
                                             moving_var_init,
-                                            use_batch_statistics,
-                                            process_groups=process_groups,
-                                            input_dims='both')
+                                            use_batch_statistics)
+        self.is_global = False
+        self.group_name = None
+        self.process_groups = process_groups
+        if self.process_groups != 0:
+            self.rank_id = get_rank()
+            self.rank_size = get_group_size()
+            if self.process_groups is not None:
+                validator.check_isinstance("process_groups", self.process_groups, list)
+                self._check_rank_ids(self.process_groups, self.rank_size)
+                self._create_sync_groups()
+            elif self.rank_size > 1:
+                self.is_global = True
+                self.group_device_num = self.rank_size
+                if context.get_context("device_target") == "Ascend":
+                    self.group_name = "hccl_world_group"
+                elif context.get_context("device_target") == "GPU":
+                    self.group_name = "nccl_world_group"
 
-    def _check_data_dim(self, x):
-        if x.dim == 0:
-            pass
+        if self.is_global:
+            self.bn_train = inner.SyncBatchNorm(epsilon=self.eps,
+                                                momentum=self.momentum,
+                                                group=self.group_name,
+                                                device_num=self.group_device_num)
+
+    def _create_sync_groups(self):
+        """ 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
+                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
+    @_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)
+        _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, 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}.")
+            seen.add(rid)
 
 
 class LayerNorm(Cell):
@@ -799,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)`,
@@ -867,7 +724,6 @@ class LayerNorm(Cell):
 
 class _InstanceNorm(Cell):
     """Instance Normalization base class."""
-
     @cell_attr_register
     def __init__(self,
                  num_features,
@@ -875,8 +731,7 @@ class _InstanceNorm(Cell):
                  momentum=0.1,
                  affine=True,
                  gamma_init='ones',
-                 beta_init='zeros',
-                 input_dims='2d'):
+                 beta_init='zeros'):
         """Initialize Normalization base class."""
         super(_InstanceNorm, self).__init__()
         validator.check_value_type('num_features', num_features, [int], self.cls_name)
@@ -893,7 +748,6 @@ class _InstanceNorm(Cell):
                              f"but got {momentum}.")
         self.num_features = num_features
         self.eps = eps
-        self.input_dims = input_dims
         self.moving_mean = Parameter(initializer('zeros', num_features), name="mean", requires_grad=False)
         self.moving_variance = Parameter(initializer('ones', num_features), name="variance", requires_grad=False)
         self.gamma = Parameter(initializer(
@@ -906,7 +760,7 @@ class _InstanceNorm(Cell):
         self.instance_bn = P.InstanceNorm(epsilon=self.eps, momentum=self.momentum)
 
     def construct(self, x):
-        _shape_check_in(self.shape(x), self.input_dims, self.cls_name)
+        self._check_input_dim(self.shape(x), self.cls_name)
         return self.instance_bn(x,
                                 self.gamma,
                                 self.beta,
@@ -952,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.
@@ -999,21 +853,12 @@ class InstanceNorm1d(_InstanceNorm):
         (2, 3, 5)
     """
 
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.1,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros'):
-        """Initialize InstanceNorm2d."""
-        super(InstanceNorm1d, self).__init__(num_features,
-                                             eps,
-                                             momentum,
-                                             affine,
-                                             gamma_init,
-                                             beta_init,
-                                             input_dims='1d')
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 3, cls_name)
+
 
 
 class InstanceNorm2d(_InstanceNorm):
@@ -1040,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.
@@ -1087,21 +932,11 @@ class InstanceNorm2d(_InstanceNorm):
         (2, 3, 2, 2)
     """
 
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.1,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros'):
-        """Initialize InstanceNorm2d."""
-        super(InstanceNorm2d, self).__init__(num_features,
-                                             eps,
-                                             momentum,
-                                             affine,
-                                             gamma_init,
-                                             beta_init,
-                                             input_dims='2d')
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 4, cls_name)
 
 
 class InstanceNorm3d(_InstanceNorm):
@@ -1128,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.
@@ -1175,21 +1010,11 @@ class InstanceNorm3d(_InstanceNorm):
         (2, 3, 5, 2, 2)
     """
 
-    def __init__(self,
-                 num_features,
-                 eps=1e-5,
-                 momentum=0.1,
-                 affine=True,
-                 gamma_init='ones',
-                 beta_init='zeros'):
-        """Initialize InstanceNorm2d."""
-        super(InstanceNorm3d, self).__init__(num_features,
-                                             eps,
-                                             momentum,
-                                             affine,
-                                             gamma_init,
-                                             beta_init,
-                                             input_dims='3d')
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 5, cls_name)
 
 
 class GroupNorm(Cell):
@@ -1212,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)` .
@@ -1273,7 +1098,7 @@ class GroupNorm(Cell):
     def _cal_output(self, x):
         """calculate groupnorm output"""
         batch, channel, height, width = self.shape(x)
-        _channel_check(channel, self.num_channels, self.cls_name)
+        self._channel_check(channel, self.num_channels, self.cls_name)
         x = self.reshape(x, (batch, self.num_groups, -1))
         mean = self.reduce_mean(x, 2)
         var = self.reduce_sum(self.square(x - mean), 2) / (channel * height * width / self.num_groups)
@@ -1283,11 +1108,32 @@ class GroupNorm(Cell):
         output = x * self.reshape(self.gamma, (-1, 1, 1)) + self.reshape(self.beta, (-1, 1, 1))
         return output
 
-    def construct(self, x):
-        _shape_check(self.shape(x), self.cls_name)
-        _check_dtype(x.dtype, [mstype.float16, mstype.float32], "input", self.cls_name)
-        output = self._cal_output(x)
-        return output
+    @staticmethod
+    @_primexpr
+    def _check_input_dim(shape, cls_name):
+        dim = len(shape)
+        _check_dim(dim, 4, cls_name)
+
+    @staticmethod
+    @_primexpr
+    def _channel_check(channel, num_channel, prim_name=None):
+        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
+    def _check_dtype(dtype, valid_dtypes, prim_name=None):
+        validator.check_type_name("input", dtype, valid_dtypes, prim_name)
 
     def extend_repr(self):
         return 'num_groups={}, num_channels={}'.format(self.num_groups, self.num_channels)
+
+    def construct(self, x):
+        self._check_input_dim(self.shape(x), self.cls_name)
+        self._check_dtype(x.dtype, [mstype.float16, mstype.float32], self.cls_name)
+        output = self._cal_output(x)
+        return output