[0-size Tensor No.175] Add 0-size Tensor support for group_norm (#73152)

co63oc · web-flow · commit 8bd5ed0a5193 · 2025-06-17T17:51:32.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/kernels/cpu/group_norm_grad_kernel.cc b/paddle/phi/kernels/cpu/group_norm_grad_kernel.cc
@@ -22,11 +22,11 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/extensions.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
-
 namespace phi {
 
 template <typename T, typename Context>
@@ -44,6 +44,29 @@ void GroupNormGradKernel(const Context& dev_ctx,
                          DenseTensor* d_x,
                          DenseTensor* d_scale,
                          DenseTensor* d_bias) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(d_x);
+    if (d_scale) {
+      // If batch dim is 0, we should set d_scale to zero, or else NAN
+      if (x.dims().size() > 0 && x.dims()[0] == 0) {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              0,
+                              d_scale);
+
+      } else {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              NAN,
+                              d_scale);
+      }
+    }
+    if (d_bias) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(d_bias->dims())), 0, d_bias);
+    }
+    return;
+  }
   const DataLayout data_layout = common::StringToDataLayout(data_layout_str);
   const auto scale_ptr = scale.get_ptr();
   const auto bias_ptr = bias.get_ptr();
diff --git a/paddle/phi/kernels/cpu/group_norm_kernel.cc b/paddle/phi/kernels/cpu/group_norm_kernel.cc
@@ -22,6 +22,7 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/extensions.h"
@@ -40,6 +41,19 @@ void GroupNormKernel(const Context& dev_ctx,
                      DenseTensor* y,
                      DenseTensor* mean,
                      DenseTensor* var) {
+  if (y && y->numel() == 0) {
+    dev_ctx.template Alloc<T>(y);
+    // mean, var are intermediate in ops yaml config.
+    if (mean) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(mean->dims())), 0, mean);
+    }
+    if (var) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(var->dims())), 0, var);
+    }
+    return;
+  }
   const DataLayout data_layout = common::StringToDataLayout(data_layout_str);
   const auto scale_ptr = scale.get_ptr();
   const auto bias_ptr = bias.get_ptr();
diff --git a/paddle/phi/kernels/gpu/group_norm_grad_kernel.cu b/paddle/phi/kernels/gpu/group_norm_grad_kernel.cu
@@ -17,9 +17,9 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/gpu/group_norm_utils.h"
-
 namespace phi {
 
 template <typename T, typename AccT, int flags>
@@ -279,6 +279,29 @@ void GroupNormGradKernel(const Context& dev_ctx,
                          DenseTensor* d_x,
                          DenseTensor* d_scale,
                          DenseTensor* d_bias) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(d_x);
+    if (d_scale) {
+      // If batch dim is 0, we should set d_scale to zero, or else NAN
+      if (x.dims().size() > 0 && x.dims()[0] == 0) {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              0,
+                              d_scale);
+
+      } else {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              NAN,
+                              d_scale);
+      }
+    }
+    if (d_bias) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(d_bias->dims())), 0, d_bias);
+    }
+    return;
+  }
   using AccT = typename phi::dtype::MPTypeTrait<T>::Type;
   const DataLayout data_layout = common::StringToDataLayout(data_layout_str);
   const auto scale_ptr = scale.get_ptr();
diff --git a/paddle/phi/kernels/gpu/group_norm_kernel.cu b/paddle/phi/kernels/gpu/group_norm_kernel.cu
@@ -23,7 +23,7 @@
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/float16.h"
 #include "paddle/phi/core/device_context.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 static inline int32_t divUp(int32_t m, int32_t n) { return (m + n - 1) / n; }
@@ -1178,6 +1178,18 @@ void GroupNormKernel(const Context& dev_ctx,
                      DenseTensor* y,
                      DenseTensor* mean,
                      DenseTensor* var) {
+  if (y && y->numel() == 0) {
+    dev_ctx.template Alloc<T>(y);
+    if (mean) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(mean->dims())), 0, mean);
+    }
+    if (var) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(var->dims())), 0, var);
+    }
+    return;
+  }
   using std::is_same;
   if (is_same<T, phi::dtype::float16>::value && data_layout_str == "NHWC") {
     const paddle::optional<DenseTensor>& residual =
diff --git a/paddle/phi/kernels/xpu/group_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/group_norm_grad_kernel.cc
@@ -22,6 +22,7 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
@@ -41,6 +42,29 @@ void GroupNormGradKernel(const Context& dev_ctx,
                          DenseTensor* d_x,
                          DenseTensor* d_scale,
                          DenseTensor* d_bias) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(d_x);
+    if (d_scale) {
+      // If batch dim is 0, we should set d_scale to zero, or else NAN
+      if (x.dims().size() > 0 && x.dims()[0] == 0) {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              0,
+                              d_scale);
+
+      } else {
+        phi::Full<T, Context>(dev_ctx,
+                              phi::IntArray(common::vectorize(d_scale->dims())),
+                              NAN,
+                              d_scale);
+      }
+    }
+    if (d_bias) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(d_bias->dims())), 0, d_bias);
+    }
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
   int ret = 0;
diff --git a/paddle/phi/kernels/xpu/group_norm_kernel.cc b/paddle/phi/kernels/xpu/group_norm_kernel.cc
@@ -22,6 +22,7 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -36,6 +37,18 @@ void GroupNormKernel(const Context& dev_ctx,
                      DenseTensor* y,
                      DenseTensor* mean,
                      DenseTensor* var) {
+  if (y && y->numel() == 0) {
+    dev_ctx.template Alloc<T>(y);
+    if (mean) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(mean->dims())), 0, mean);
+    }
+    if (var) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(var->dims())), 0, var);
+    }
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
 
   const DataLayout data_layout = common::StringToDataLayout(data_layout_str);
diff --git a/test/legacy_test/test_group_norm_op_v2.py b/test/legacy_test/test_group_norm_op_v2.py
@@ -35,7 +35,11 @@ def group_norm_naive_for_general_dimension(
     input_shape = x.shape
     N, C = x.shape[0], x.shape[1]
     G = groups
-    x = x.reshape((N * G, -1))
+    if 0 in x.shape:
+        # output will reshape to input_shape
+        x = x.reshape((N * G, 0))
+    else:
+        x = x.reshape((N * G, -1))
     mean = np.mean(x, axis=1, keepdims=True)
     var = np.var(x, axis=1, keepdims=True)
     output = (x - mean) / np.sqrt(var + epsilon)
@@ -594,6 +598,76 @@ def test_numerical_accuracy(self):
             )
 
 
+class TestGroupNormAPIV2_ZeroSize(unittest.TestCase):
+    def test_numerical_accuracy(self):
+        paddle.disable_static()
+        shape_groups = [
+            [(2, 4, 3, 2, 0), 4],
+            [(0, 1, 0, 0, 1), 1],
+        ]
+        np.random.seed(10)
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not core.is_compiled_with_cuda()
+        ):
+            places.append(base.CPUPlace())
+        if core.is_compiled_with_cuda() and core.op_support_gpu("group_norm"):
+            places.append(base.CUDAPlace(0))
+
+        for place in places:
+            paddle.disable_static(place)
+            for shape_group in shape_groups:
+                shape = shape_group[0]
+                group = shape_group[1]
+                scale = np.array([1]).astype("float32")
+                bias = np.array([0]).astype("float32")
+                data = np.random.random(shape).astype("float32")
+                expect_res1 = group_norm_naive_for_general_dimension(
+                    data,
+                    scale,
+                    bias,
+                    epsilon=1e-5,
+                    groups=group,
+                )
+
+                data_pd = paddle.to_tensor(data)
+                data_pd.stop_gradient = False
+                scale_ = paddle.ones(
+                    shape[1]
+                )  # shape[1] is the number of channels
+                scale_.stop_gradient = False
+                bias_ = paddle.zeros(shape[1])
+                bias_.stop_gradient = False
+                result1 = paddle.nn.functional.group_norm(
+                    data_pd,
+                    group,
+                    data_format='NCDHW',
+                    weight=scale_,
+                    bias=bias_,
+                )
+                np.testing.assert_allclose(
+                    result1.numpy(), expect_res1, atol=1e-5
+                )
+
+                loss = paddle.sum(result1)
+                loss.backward()
+                np.testing.assert_allclose(
+                    data_pd.grad.shape,
+                    data_pd.shape,
+                )
+                # If batch is 0, scale grad is 0, or else nan.
+                if data_pd.shape[0] == 0:
+                    scale2 = paddle.zeros(scale_.shape)
+                else:
+                    scale2 = paddle.full(scale_.shape, paddle.nan)
+                np.testing.assert_allclose(scale_.grad.numpy(), scale2.numpy())
+                np.testing.assert_allclose(
+                    bias_.grad.numpy(), paddle.zeros(bias_.shape).numpy()
+                )
+
+
 class TestGroupNormDimException(unittest.TestCase):
     def test_exception(self):
         with paddle.pir_utils.OldIrGuard():
