API improvement nn.functional.group_norm 易用性提升 (#62672)

NKNaN · web-flow · commit d7768a77817f · 2024-03-25T11:21:06.000+08:00
* add nn.functional.group_norm

* fix docs

* fix docs
diff --git a/python/paddle/nn/functional/__init__.py b/python/paddle/nn/functional/__init__.py
@@ -122,6 +122,7 @@
 )
 from .norm import (
     batch_norm,
+    group_norm,
     instance_norm,
     layer_norm,
     local_response_norm,
@@ -276,4 +277,5 @@
     'soft_margin_loss',
     'gaussian_nll_loss',
     'scaled_dot_product_attention',
+    'group_norm',
 ]
diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py
@@ -637,3 +637,116 @@ def local_response_norm(
     div = paddle.pow(div, beta)
     res = paddle.divide(x, div, name=name)
     return res
+
+
+def group_norm(
+    x,
+    num_groups,
+    epsilon=1e-05,
+    weight=None,
+    bias=None,
+    data_format='NCHW',
+    name=None,
+):
+    """
+    nn.GroupNorm is recommended.
+    For more information, please refer to :ref:`api_paddle_nn_GroupNorm` .
+
+    Parameters:
+        x(Tensor): Input Tensor with shape: attr:`(batch, num_features, *)`.
+        num_groups(int): The number of groups that divided from channels.
+        epsilon(float, optional): The small value added to the variance to prevent
+            division by zero. Default: 1e-05.
+        weight(Tensor, optional): The weight Tensor of group_norm, with shape: attr:`[num_channels]`.
+            Default: None.
+        bias(Tensor, optional): The bias Tensor of group_norm, with shape: attr:`[num_channels]`.
+            Default: None.
+        data_format(str, optional): Specify the input data format. Only NCHW is supported. Default: NCHW.
+        name(str, optional): Name for the GroupNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
+
+    Returns:
+        Tensor, the output has the same shape with ``x``.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> paddle.seed(100)
+            >>> x = paddle.arange(48, dtype="float32").reshape((2, 6, 2, 2))
+            >>> group_norm_out = paddle.nn.functional.group_norm(x, num_groups=6)
+
+            >>> print(group_norm_out)
+            Tensor(shape=[2, 6, 2, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[[[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]]],
+             [[[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]],
+              [[-1.34163547, -0.44721183],
+               [ 0.44721183,  1.34163547]]]])
+    """
+    if data_format not in ['NCHW', 'NHWC']:
+        raise ValueError("unsupported data layout:" + data_format)
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.group_norm(
+            x,
+            weight,
+            bias,
+            epsilon,
+            num_groups,
+            data_format,
+        )
+    else:
+        helper = LayerHelper('group_norm', **locals())
+        mean_out = helper.create_variable_for_type_inference(
+            dtype=x.dtype, stop_gradient=True
+        )
+        variance_out = helper.create_variable_for_type_inference(
+            dtype=x.dtype, stop_gradient=True
+        )
+
+        inputs = {'X': x}
+        if bias is not None:
+            inputs['Bias'] = bias
+        if weight is not None:
+            inputs['Scale'] = weight
+
+        # create output
+        group_norm_out = helper.create_variable_for_type_inference(
+            dtype=x.dtype
+        )
+
+        helper.append_op(
+            type="group_norm",
+            inputs=inputs,
+            outputs={
+                "Y": group_norm_out,
+                "Mean": mean_out,
+                "Variance": variance_out,
+            },
+            attrs={
+                "epsilon": epsilon,
+                "groups": num_groups,
+                "data_layout": data_format,
+            },
+        )
+
+        return helper.append_activation(group_norm_out)
diff --git a/python/paddle/nn/layer/norm.py b/python/paddle/nn/layer/norm.py
@@ -46,7 +46,7 @@
     no_grad,
 )
 from .. import functional as F
-from ..functional import batch_norm, instance_norm, layer_norm
+from ..functional import batch_norm, group_norm, instance_norm, layer_norm
 from ..initializer import Constant, Normal
 from .layers import Layer
 
@@ -533,51 +533,15 @@ def __init__(
             )
 
     def forward(self, input):
-        if in_dynamic_or_pir_mode():
-            return _C_ops.group_norm(
-                input,
-                self.weight,
-                self.bias,
-                self._epsilon,
-                self._num_groups,
-                self._data_format,
-            )
-
-        mean_out = self._helper.create_variable_for_type_inference(
-            dtype=input.dtype, stop_gradient=True
-        )
-        variance_out = self._helper.create_variable_for_type_inference(
-            dtype=input.dtype, stop_gradient=True
-        )
-
-        inputs = {'X': input}
-        if self.bias is not None:
-            inputs['Bias'] = self.bias
-        if self.weight is not None:
-            inputs['Scale'] = self.weight
-
-        # create output
-        group_norm_out = self._helper.create_variable_for_type_inference(
-            dtype=input.dtype
-        )
-
-        self._helper.append_op(
-            type="group_norm",
-            inputs=inputs,
-            outputs={
-                "Y": group_norm_out,
-                "Mean": mean_out,
-                "Variance": variance_out,
-            },
-            attrs={
-                "epsilon": self._epsilon,
-                "groups": self._num_groups,
-                "data_layout": self._data_format,
-            },
+        return group_norm(
+            input,
+            self._num_groups,
+            self._epsilon,
+            self.weight,
+            self.bias,
+            self._data_format,
         )
 
-        return self._helper.append_activation(group_norm_out, None)
-
     def extra_repr(self):
         return 'num_groups={}, num_channels={}, epsilon={}'.format(
             self._num_groups, self._num_channels, self._epsilon

Original file line number	Diff line number	Diff line change
`@@ -122,6 +122,7 @@`
`122`	`122`	`)`
`123`	`123`	`from .norm import (`
`124`	`124`	`batch_norm,`
	`125`	`+ group_norm,`
`125`	`126`	`instance_norm,`
`126`	`127`	`layer_norm,`
`127`	`128`	`local_response_norm,`
`@@ -276,4 +277,5 @@`
`276`	`277`	`'soft_margin_loss',`
`277`	`278`	`'gaussian_nll_loss',`
`278`	`279`	`'scaled_dot_product_attention',`
	`280`	`+ 'group_norm',`
`279`	`281`	`]`