mindspore 2.2.14__cp37-cp37m-manylinux1_x86_64.whl → 2.3.0rc2__cp37-cp37m-manylinux1_x86_64.whl

mindspore/nn/layer/embedding.py

@@ -40,7 +40,8 @@ __all__ = ['Embedding', 'EmbeddingLookup', 'MultiFieldEmbeddingLookup']
 @_primexpr
 def _check_input_2d(input_shape, param_name, func_name):
     if len(input_shape) != 2:
-        raise ValueError(f"For '{func_name}', the dimension of '{param_name}' must be 2d, but got {len(input_shape)}")
+        raise ValueError(
+            f"For '{func_name}', the dimension of '{param_name}' must be 2d, but got {len(input_shape)}")
 
 
 @constexpr
@@ -65,7 +66,7 @@ class Embedding(Cell):
         use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: ``False`` .
         embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
             Refer to class `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             for the values of string when a string is specified. Default: ``'normal'`` .
         dtype (:class:`mindspore.dtype`): Data type of `x`. Default: ``mstype.float32`` .
         padding_idx (int, None): When the padding_idx encounters index, the output embedding vector of this index
@@ -104,23 +105,33 @@ class Embedding(Cell):
                  dtype=mstype.float32, padding_idx=None):
         """Initialize Embedding."""
         super(Embedding, self).__init__()
-        self.vocab_size = Validator.check_value_type('vocab_size', vocab_size, [int], self.cls_name)
-        self.embedding_size = Validator.check_value_type('embedding_size', embedding_size, [int], self.cls_name)
-        Validator.check_value_type('use_one_hot', use_one_hot, [bool], self.cls_name)
-        Validator.check_subclass("dtype", dtype, mstype.number_type, self.cls_name)
+        self.vocab_size = Validator.check_value_type(
+            'vocab_size', vocab_size, [int], self.cls_name)
+        self.embedding_size = Validator.check_value_type(
+            'embedding_size', embedding_size, [int], self.cls_name)
+        Validator.check_value_type(
+            'use_one_hot', use_one_hot, [bool], self.cls_name)
+        Validator.check_subclass(
+            "dtype", dtype, mstype.number_type, self.cls_name)
         self.use_one_hot = use_one_hot
         self.dtype = dtype
-        self.init_tensor = initializer(embedding_table, [vocab_size, embedding_size])
+        self.init_tensor = initializer(
+            embedding_table, [vocab_size, embedding_size])
         self.padding_idx = padding_idx
         if padding_idx is not None:
             self.padding_idx = Validator.check_int_range(padding_idx, 0, vocab_size, Validator.INC_LEFT,
                                                          "padding_idx", self.cls_name)
             if isinstance(self.init_tensor, Tensor) and self.init_tensor.init is not None:
                 self.init_tensor = self.init_tensor.init_data()
-            self.init_tensor = self.init_tensor.asnumpy()
+            init_tensor_type = self.init_tensor.dtype
+            if init_tensor_type == mstype.bfloat16:
+                self.init_tensor = self.init_tensor.float().asnumpy()
+            else:
+                self.init_tensor = self.init_tensor.asnumpy()
             self.init_tensor[self.padding_idx] = 0
-            self.init_tensor = Tensor(self.init_tensor)
-        self.embedding_table = Parameter(self.init_tensor, name='embedding_table')
+            self.init_tensor = Tensor(self.init_tensor, init_tensor_type)
+        self.embedding_table = Parameter(
+            self.init_tensor, name='embedding_table')
         self.expand = P.ExpandDims()
         self.reshape_flat = P.Reshape()
         self.shp_flat = (-1,)
@@ -138,8 +149,10 @@ class Embedding(Cell):
         flat_ids = self.reshape_flat(ids, self.shp_flat)
 
         if self.use_one_hot:
-            one_hot_ids = self.one_hot(flat_ids, self.vocab_size, self.on_value, self.off_value)
-            output_for_reshape = self.array_mul(one_hot_ids, self.embedding_table)
+            one_hot_ids = self.one_hot(
+                flat_ids, self.vocab_size, self.on_value, self.off_value)
+            output_for_reshape = self.array_mul(
+                one_hot_ids, self.embedding_table)
         else:
             output_for_reshape = self.gather(self.embedding_table, flat_ids, 0)
 
@@ -181,8 +194,18 @@ class EmbeddingLookup(Cell):
             is specified. Default: ``'normal'`` .
         target (str): Specifies the target where the op is executed. The value must in
             [ ``'DEVICE'`` , ``'CPU'`` ]. Default: ``'CPU'`` .
-        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through
-            :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'`` .
+        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. Default: ``'batch_slice'`` .
+
+          - batch_slice (str): Divides the input index tensor into batches and retrieves
+            the corresponding embedding vectors. This is applicable when each sample has the same number of indices.
+          - field_slice (str): Divides the input index tensor into fields and retrieves the corresponding embedding
+            vectors. This is applicable when each sample may have a different number of indices, but have the same
+            feature dimensions.
+          - table_row_slice (str): Treats the input index tensor as a 2D table, divides it by rows, and retrieves
+            the corresponding embedding vectors.
+          - table_column_slice (str): Treats the input index tensor as a 2D table, divides it by columns, and retrieves
+            the corresponding embedding vectors.
+
         manual_shapes (tuple): The accompaniment array in field slice mode. Default: ``None`` .
         max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32
                                        or None. Default: ``None`` .
@@ -237,13 +260,16 @@ class EmbeddingLookup(Cell):
         """Initialize EmbeddingLookup."""
         super(EmbeddingLookup, self).__init__()
         Validator.check_value_type('sparse', sparse, [bool], self.cls_name)
-        self.vocab_size = Validator.check_positive_int(vocab_size, 'vocab_size')
-        self.vocab_cache_size = Validator.check_non_negative_int(vocab_cache_size, 'vocab_cache_size')
+        self.vocab_size = Validator.check_positive_int(
+            vocab_size, 'vocab_size')
+        self.vocab_cache_size = Validator.check_non_negative_int(
+            vocab_cache_size, 'vocab_cache_size')
         self.target = target
         self.sparse = sparse
         self.cache_enable = self.vocab_cache_size > 0
         self.forward_unique = False
-        Validator.check_string(target, ['CPU', 'DEVICE'], 'target', self.cls_name)
+        Validator.check_string(
+            target, ['CPU', 'DEVICE'], 'target', self.cls_name)
         if not sparse and target == 'CPU':
             raise ValueError(f"For '{self.cls_name}', 'sparse' must be True when 'target' is \"CPU\", "
                              f"but got 'sparse': {sparse} and 'target': {target}")
@@ -256,11 +282,13 @@ class EmbeddingLookup(Cell):
         enable_ps = _get_ps_context("enable_ps")
         if enable_ps:
             self._process_vocab_cache(slice_mode)
-        self.embedding_size = Validator.check_positive_int(embedding_size, 'embedding_size', self.cls_name)
+        self.embedding_size = Validator.check_positive_int(
+            embedding_size, 'embedding_size', self.cls_name)
         self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size],
                                                      dtype=dtype), name='embedding_table')
         parallel_mode = _get_parallel_mode()
-        is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
+        is_auto_parallel = parallel_mode in (
+            ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
         self.gather_revert = P.Gather()
         self.reshape_first = P.Reshape()
         self.reshape = P.Reshape()
@@ -269,7 +297,8 @@ class EmbeddingLookup(Cell):
         if is_auto_parallel:
             self.unique = P.Unique().shard(((1,),))
         if self.cache_enable and enable_ps:
-            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init, dtype=dtype)
+            self._set_voacb_cache_enable_for_ps(
+                vocab_cache_size, embedding_size, vocab_size, param_init, dtype=dtype)
             if is_auto_parallel:
                 self.unique.add_prim_attr('cache_enable', True)
         indices_shape_size = 2
@@ -281,11 +310,13 @@ class EmbeddingLookup(Cell):
                 raise TypeError(f"For '{self.cls_name}', the type of 'manual_shapes' must be tuple(int), "
                                 f"but got {type(manual_shapes).__name__}!")
             for dim in manual_shapes:
-                Validator.check_positive_int(dim, 'manual shape dim', self.cls_name)
+                Validator.check_positive_int(
+                    dim, 'manual shape dim', self.cls_name)
             self.gatherv2.add_prim_attr("manual_split", manual_shapes)
             self.embeddinglookup.add_prim_attr("manual_split", manual_shapes)
             self.gatherv2.shard(((get_group_size(), 1), (1, get_group_size())))
-            self.embeddinglookup.shard(((get_group_size(), 1), (1, get_group_size())))
+            self.embeddinglookup.shard(
+                ((get_group_size(), 1), (1, get_group_size())))
         elif slice_mode == "table_row_slice" and is_auto_parallel:
             full_batch = _get_full_batch()
             if (target == 'DEVICE' and not full_batch) or (self.cache_enable and enable_ps and sparse):
@@ -294,7 +325,8 @@ class EmbeddingLookup(Cell):
                 self.forward_unique = True
             indices_strategy = (1,)*indices_shape_size
             self.gatherv2.shard(((get_group_size(), 1), indices_strategy))
-            self.embeddinglookup.shard(((get_group_size(), 1), indices_strategy))
+            self.embeddinglookup.shard(
+                ((get_group_size(), 1), indices_strategy))
         elif slice_mode == "table_column_slice" and is_auto_parallel:
             if target == 'DEVICE':
                 indices_shape_size = 1
@@ -302,7 +334,8 @@ class EmbeddingLookup(Cell):
                 self.forward_unique = True
             indices_strategy = (1,)*indices_shape_size
             self.gatherv2.shard(((1, get_group_size()), indices_strategy))
-            self.embeddinglookup.shard(((1, get_group_size()), indices_strategy))
+            self.embeddinglookup.shard(
+                ((1, get_group_size()), indices_strategy))
         elif slice_mode == "batch_slice" and is_auto_parallel:
             indices_strategy = [get_group_size()]
             indices_strategy.extend([1] * (indices_shape_size - 1))
@@ -311,15 +344,18 @@ class EmbeddingLookup(Cell):
             self.embeddinglookup.shard(((1, 1), indices_strategy))
         else:
             if is_auto_parallel:
-                support_mode = ["field_slice", "table_row_slice", "table_column_slice", "batch_slice"]
+                support_mode = ["field_slice", "table_row_slice",
+                                "table_column_slice", "batch_slice"]
                 raise ValueError(f"For '{self.cls_name}', the 'slice_mode' must be in {support_mode}, "
                                  f"but got \"{slice_mode}\".")
         if self.cache_enable and not enable_ps:
-            raise ValueError(f"For '{self.cls_name}', haven't supported cache enable for not ps mode.")
+            raise ValueError(
+                f"For '{self.cls_name}', haven't supported cache enable for not ps mode.")
         self.embedding_table.unique = self.forward_unique
         self.max_norm = max_norm
         if self.max_norm is not None:
-            self.max_norm = Validator.check_positive_float(self.max_norm, 'max_norm', self.cls_name)
+            self.max_norm = Validator.check_positive_float(
+                self.max_norm, 'max_norm', self.cls_name)
             self.max_norm = Tensor(self.max_norm, dtype=mstype.float32)
 
     def _process_vocab_cache(self, slice_mode):
@@ -337,7 +373,8 @@ class EmbeddingLookup(Cell):
                 return
             self.is_ps_server = _is_role_pserver() and _enable_distributed_mindrt()
             parallel_mode = _get_parallel_mode()
-            is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
+            is_auto_parallel = parallel_mode in (
+                ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
             if is_auto_parallel:
                 rank_size = get_group_size()
                 rank_id = get_rank()
@@ -366,7 +403,8 @@ class EmbeddingLookup(Cell):
             self.embedding_table.is_param_ps = True
             self.embedding_table.cache_enable = True
             self.embedding_table.key = param_key
-            _insert_hash_table_size(self.embedding_table.name, vocab_cache_size, embedding_size, vocab_size, param_key)
+            _insert_hash_table_size(
+                self.embedding_table.name, vocab_cache_size, embedding_size, vocab_size, param_key)
 
         if _enable_distributed_mindrt():
             self.rank_id = get_rank()
@@ -392,7 +430,8 @@ class EmbeddingLookup(Cell):
             raise ValueError("The Parameter Server number is zero.")
         # Assign the embedding table dimensions.
         for _ in range(server_num):
-            self.embedding_table_vocab_dim_list.append(self.vocab_size // server_num)
+            self.embedding_table_vocab_dim_list.append(
+                self.vocab_size // server_num)
         rest_vocab_size = self.vocab_size % server_num
         if rest_vocab_size != 0:
             for i in range(rest_vocab_size):
@@ -416,7 +455,8 @@ class EmbeddingLookup(Cell):
                     embedding_lookup = P.SparseGatherV2()
                 else:
                     embedding_lookup = P.Gather()
-                embedding_lookup.add_prim_attr('offset', self.embedding_offset[i])
+                embedding_lookup.add_prim_attr(
+                    'offset', self.embedding_offset[i])
             embedding_lookup.add_prim_attr('rank_id', i)
             embedding_lookup.add_prim_attr('ms_role', 'MS_PSERVER')
             self.embedding_lookup_list.append(embedding_lookup)
@@ -466,8 +506,10 @@ class EmbeddingLookup(Cell):
                 shp = self.shape(indices) + (self.embedding_size,)
                 indices_flatten = self.reshape_first(indices, (-1,))
                 unique_id, unique_idx = self.unique(indices_flatten)
-                weight_unique = self.gatherv2(self.embedding_table, unique_id, 0)
-                weight_flatten = self.gather_revert(weight_unique, unique_idx, 0)
+                weight_unique = self.gatherv2(
+                    self.embedding_table, unique_id, 0)
+                weight_flatten = self.gather_revert(
+                    weight_unique, unique_idx, 0)
                 out = self.reshape(weight_flatten, shp)
             else:
                 out = self.gatherv2(self.embedding_table, indices, 0)
@@ -504,8 +546,18 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
             is specified. Default: ``'normal'`` .
         target (str): Specifies the target where the op is executed. The value must in
             [ ``'DEVICE'`` , ``'CPU'`` ]. Default: ``'CPU'`` .
-        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through
-            :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'``.
+        slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. Default: ``'batch_slice'``.
+
+          - batch_slice (str): Divides the input index tensor into batches and retrieves
+            the corresponding embedding vectors. This is applicable when each sample has the same number of indices.
+          - field_slice (str): Divides the input index tensor into fields and retrieves the corresponding embedding
+            vectors. This is applicable when each sample may have a different number of indices, but have the same
+            feature dimensions.
+          - table_row_slice (str): Treats the input index tensor as a 2D table, divides it by rows, and retrieves
+            the corresponding embedding vectors.
+          - table_column_slice (str): Treats the input index tensor as a 2D table, divides it by columns, and retrieves
+            the corresponding embedding vectors.
+
         feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently.
             Default:  ``None`` .
         max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32.
@@ -564,7 +616,8 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         """Initialize MultiFieldEmbeddingLookup."""
         super(MultiFieldEmbeddingLookup, self).__init__(vocab_size, embedding_size, param_init, target,
                                                         slice_mode, feature_num_list, max_norm, sparse, dtype=dtype)
-        self.field_size = Validator.check_positive_int(field_size, 'field_size', self.cls_name)
+        self.field_size = Validator.check_positive_int(
+            field_size, 'field_size', self.cls_name)
         self.operator = operator
 
         self.mul = P.Mul()
@@ -582,7 +635,8 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         self.max_mask_mul = P.Mul()
         self.max_no_equal = P.NotEqual()
 
-        Validator.check_string(operator, ['SUM', 'MAX', 'MEAN'], 'operator', self.cls_name)
+        Validator.check_string(
+            operator, ['SUM', 'MAX', 'MEAN'], 'operator', self.cls_name)
         if operator == MultiFieldEmbeddingLookup.OPERATOR_SUM:
             self.merge_op = P.UnsortedSegmentSum()
         elif operator == MultiFieldEmbeddingLookup.OPERATOR_MAX:
@@ -665,13 +719,16 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
             clip_by_norm = ClipByNorm(axis)
             out = clip_by_norm(out, self.max_norm)
 
-        weights = self.reshape(input_values, (batch_size, self.shape(input_indices)[1], 1))
+        weights = self.reshape(
+            input_values, (batch_size, self.shape(input_indices)[1], 1))
         embedding = self.mul(weights, out)
 
         if self.operator == 'MAX':
             # Fill the padding value to -inf, so the padded value will not influence the results
-            negative_inf_mask = self.cast(self.equal(weights, 0), mstype.float32)
-            inf_mask = self.inf_mask_mul(negative_inf_mask, self.negative_inf_value)
+            negative_inf_mask = self.cast(
+                self.equal(weights, 0), mstype.float32)
+            inf_mask = self.inf_mask_mul(
+                negative_inf_mask, self.negative_inf_value)
             embedding = self.inf_add(embedding, inf_mask)
             embedding = self.reshape(embedding, (-1, self.embedding_size))
             field_ids = self.reshape(field_ids, (-1,))
@@ -679,15 +736,19 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         merged_vectors = self.merge_op(embedding, field_ids, num_segments)
 
         if self.operator == 'MAX':
-            value_count = self.count_op(self.abs(self.reshape(input_values, (-1,))), field_ids, num_segments)
-            value_zeros = self.cast(self.max_no_equal(value_count, 0.0), mstype.float32)
+            value_count = self.count_op(self.abs(self.reshape(
+                input_values, (-1,))), field_ids, num_segments)
+            value_zeros = self.cast(self.max_no_equal(
+                value_count, 0.0), mstype.float32)
             count = self.expand(value_zeros, -1)
             merged_vectors = self.max_mask_mul(merged_vectors, count)
 
         if self.operator == 'MEAN':
-            value_count = self.count_op(self.abs(input_values), field_ids, num_segments)
+            value_count = self.count_op(
+                self.abs(input_values), field_ids, num_segments)
             value_count = self.expand(value_count, -1)
             merged_vectors = self.div_no_nan(merged_vectors, value_count)
 
-        merged_vectors = self.reshape(merged_vectors, (batch_size, self.field_size, -1))
+        merged_vectors = self.reshape(
+            merged_vectors, (batch_size, self.field_size, -1))
         return merged_vectors

mindspore/nn/layer/image.py

@@ -115,7 +115,7 @@ def _get_dtype_max(dtype):
     """get max of the dtype"""
     np_type = mstype.dtype_to_nptype(dtype)
     if issubclass(np_type, numbers.Integral):
-        dtype_max = np.float64(np.iinfo(np_type).max)
+        dtype_max = np.float64(np.iinfo(np_type).max).item()
     else:
         dtype_max = 1.0
     return dtype_max
@@ -283,8 +283,7 @@ class SSIM(Cell):
         _check_input_filter_size(F.shape(img1), "img1", self.filter_size, self.cls_name)
         inner.SameTypeShape()(img1, img2)
         dtype_max_val = _get_dtype_max(F.dtype(img1))
-        max_val = F.scalar_cast(self.max_val, F.dtype(img1))
-        max_val = _convert_img_dtype_to_float32(max_val, dtype_max_val)
+        max_val = _convert_img_dtype_to_float32(self.max_val, dtype_max_val)
         img1 = _convert_img_dtype_to_float32(img1, dtype_max_val)
         img2 = _convert_img_dtype_to_float32(img2, dtype_max_val)
 
@@ -391,8 +390,7 @@ class MSSSIM(Cell):
         _check_input_dtype(F.dtype(img1), 'img1', valid_type, self.cls_name)
         inner.SameTypeShape()(img1, img2)
         dtype_max_val = _get_dtype_max(F.dtype(img1))
-        max_val = F.scalar_cast(self.max_val, F.dtype(img1))
-        max_val = _convert_img_dtype_to_float32(max_val, dtype_max_val)
+        max_val = _convert_img_dtype_to_float32(self.max_val, dtype_max_val)
         img1 = _convert_img_dtype_to_float32(img1, dtype_max_val)
         img2 = _convert_img_dtype_to_float32(img2, dtype_max_val)
 
@@ -469,8 +467,7 @@ class PSNR(Cell):
         _check_input_4d(F.shape(img2), "img2", self.cls_name)
         inner.SameTypeShape()(img1, img2)
         dtype_max_val = _get_dtype_max(F.dtype(img1))
-        max_val = F.scalar_cast(self.max_val, F.dtype(img1))
-        max_val = _convert_img_dtype_to_float32(max_val, dtype_max_val)
+        max_val = _convert_img_dtype_to_float32(self.max_val, dtype_max_val)
         img1 = _convert_img_dtype_to_float32(img1, dtype_max_val)
         img2 = _convert_img_dtype_to_float32(img2, dtype_max_val)
 

mindspore/nn/layer/normalization.py

@@ -21,7 +21,6 @@ import numbers
 import hashlib
 
 from mindspore.ops import operations as P
-from mindspore.ops import functional as F
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common.parameter import Parameter
 from mindspore.common.initializer import initializer, Initializer
@@ -36,6 +35,7 @@ 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
+from mindspore.ops import group_norm
 
 __all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LayerNorm', 'GroupNorm',
            'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d']
@@ -128,15 +128,19 @@ class _BatchNorm(Cell):
         self.assign_sub_mean = P.AssignSub().shard(data_parallel_strategy)
         self.assign_sub_var = P.AssignSub().shard(data_parallel_strategy)
 
-
     @staticmethod
     @_primexpr
     def _check_input_dim(shape, cls_name):
         raise NotImplementedError
 
-
     def construct(self, x):
         self._check_input_dim(self.shape(x), self.cls_name)
+        x_shape = self.shape(x)
+        reshaped_x = x
+        if len(x_shape) == 2:
+            reshaped_x = self.reshape(x, (x_shape[0], x_shape[1], 1, 1))
+        elif len(x_shape) == 3:
+            reshaped_x = self.reshape(x, (x_shape[0], x_shape[1], x_shape[2], 1))
         if self.use_batch_statistics is None:
             if self.training:
                 return self.bn_train(x,
@@ -145,11 +149,14 @@ class _BatchNorm(Cell):
                                      self.moving_mean,
                                      self.moving_variance)[0]
             if not self.training:
-                return self.bn_infer(x,
-                                     self.gamma,
-                                     self.beta,
-                                     self.moving_mean,
-                                     self.moving_variance)[0]
+                bn_out = self.bn_infer(reshaped_x,
+                                       self.gamma,
+                                       self.beta,
+                                       self.moving_mean,
+                                       self.moving_variance)[0]
+                if len(x_shape) < 4:
+                    bn_out = self.reshape(bn_out, x_shape)
+                return bn_out
 
         if self.use_batch_statistics:
             return self.bn_train(x,
@@ -158,11 +165,14 @@ class _BatchNorm(Cell):
                                  self.moving_mean,
                                  self.moving_variance)[0]
 
-        return self.bn_infer(x,
-                             self.gamma,
-                             self.beta,
-                             self.moving_mean,
-                             self.moving_variance)[0]
+        bn_out = self.bn_infer(reshaped_x,
+                               self.gamma,
+                               self.beta,
+                               self.moving_mean,
+                               self.moving_variance)[0]
+        if len(x_shape) < 4:
+            bn_out = self.reshape(bn_out, x_shape)
+        return bn_out
 
     def extend_repr(self):
         return 'num_features={}, eps={}, momentum={}, gamma={}, beta={}, moving_mean={}, moving_variance={}'.format(
@@ -196,19 +206,19 @@ class BatchNorm1d(_BatchNorm):
             Default: ``True`` .
         gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             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 `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'zeros'`` .
         moving_var_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving variance.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         use_batch_statistics (bool): If ``true`` , use the mean value and variance value of current batch data. If
             ``false`` , use the mean value and variance value of specified value. If ``None`` , the training process
@@ -292,19 +302,19 @@ class BatchNorm2d(_BatchNorm):
             Default: ``True`` .
         gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             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 `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'zeros'`` .
         moving_var_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving variance.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         use_batch_statistics (bool): Default: ``None`` .
 
@@ -381,19 +391,19 @@ class BatchNorm3d(Cell):
         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.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the beta weight.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             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 `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'zeros'`` .
         moving_var_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the moving variance.
             The values of str refer to the function `mindspore.common.initializer
-            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            <https://www.mindspore.cn/docs/en/master/api_python/mindspore.common.initializer.html>`_
             including ``'zeros'`` , ``'ones'`` , etc. Default: ``'ones'`` .
         use_batch_statistics (bool): If true, use the mean value and variance value of current batch data. If
             ``false``, use the mean value and variance value of specified value. If ``None`` , the training process
@@ -460,7 +470,6 @@ class BatchNorm3d(Cell):
         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)
@@ -549,14 +558,14 @@ 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.2/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/rank_table.html>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/mpirun.html>`_ .
 
             For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -649,16 +658,6 @@ class SyncBatchNorm(_BatchNorm):
                     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):
@@ -668,6 +667,15 @@ class SyncBatchNorm(_BatchNorm):
                                  f"but got {process_groups}.")
             seen.add(rid)
 
+    @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)
+
 
 class LayerNorm(Cell):
     r"""
@@ -1092,7 +1100,8 @@ class GroupNorm(Cell):
         dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
-        - **x** (Tensor) - The input feature with shape :math:`(N, C, H, W)` .
+        - **x** (Tensor) - The input feature with shape :math:`(N, C, *)`, where :math:`*` means, any number of
+          additional dimensions.
 
     Outputs:
         Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `x`.
@@ -1137,34 +1146,13 @@ class GroupNorm(Cell):
         self.affine = validator.check_bool(affine, arg_name="affine", prim_name=self.cls_name)
 
         self.gamma = Parameter(initializer(
-            gamma_init, num_channels, dtype=dtype), name="gamma", requires_grad=affine)
+            gamma_init, self.num_channels, dtype=dtype), name="gamma", requires_grad=affine)
         self.beta = Parameter(initializer(
-            beta_init, num_channels, dtype=dtype), name="beta", requires_grad=affine)
-        self.reduce_mean = P.ReduceMean(keep_dims=True)
-        self.reduce_sum = P.ReduceSum(keep_dims=True)
-        self.shape = F.shape
-        self.reshape = F.reshape
-        self.square = F.square
-        self.sqrt = P.Sqrt()
+            beta_init, self.num_channels, dtype=dtype), name="beta", requires_grad=affine)
 
     def _cal_output(self, x):
         """calculate groupnorm output"""
-        batch, channel, height, width = F.shape(x)
-        self._channel_check(channel, self.num_channels, self.cls_name)
-        x = F.reshape(x, (batch, self.num_groups, -1))
-        mean = self.reduce_mean(x, 2)
-        var = F.div(self.reduce_sum(F.square(F.sub(x, mean)), 2), (channel * height * width / self.num_groups))
-        std = self.sqrt(var + self.eps)
-        x = F.div(F.sub(x, mean), std)
-        x = F.reshape(x, (batch, channel, height, width))
-        output = F.add(x * F.reshape(self.gamma, (-1, 1, 1)), F.reshape(self.beta, (-1, 1, 1)))
-        return output
-
-    @staticmethod
-    @_primexpr
-    def _check_input_dim(shape, cls_name):
-        dim = len(shape)
-        _check_dim(dim, 4, cls_name)
+        return group_norm(x, self.num_groups, self.gamma, self.beta, self.eps)
 
     @staticmethod
     @_primexpr
@@ -1185,7 +1173,5 @@ class GroupNorm(Cell):
         return 'num_groups={}, num_channels={}'.format(self.num_groups, self.num_channels)
 
     def construct(self, x):
-        self._check_input_dim(F.shape(x), self.cls_name)
-        self._check_dtype(x.dtype, [mstype.float16, mstype.float32], self.cls_name)
         output = self._cal_output(x)
         return output