[Prim][PIR] Support the dynamic shape for batch_norm_grad (#67790)

zeroRains · web-flow · commit 3d65d41fcd3b · 2024-09-06T16:27:17.000+08:00
* support the dynamic shape for batch_norm_grad

* fix the dynamic shape checking
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1884,10 +1884,19 @@ void batch_norm_grad(const Tensor& x,
                      Tensor* bias_grad) {
   use_global_stats = is_test || use_global_stats;
 
+  bool has_dynamic_shape_for_x = has_dynamic_shape(x.shape());
+
   DataLayout data_layout_ = common::StringToDataLayout(data_layout);
 
   Tensor x_data = x;
   Tensor out_grad_data = out_grad;
+  auto run_var = variance_out.get();
+  auto run_mean = mean_out.get();
+  Tensor mean_data;
+  Tensor rsqrt_var;
+  std::vector<int> nchw_to_nhwc_dim = {0, 2, 3, 1};
+  std::vector<int> nhwc_to_nchw_dim = {0, 3, 1, 2};
+  auto reduce_axis = IntArray(std::vector<int64_t>{0, 1, 2});
 
   bool need_cast = x.dtype() == phi::DataType::FLOAT16 ||
                    x.dtype() == phi::DataType::BFLOAT16;
@@ -1898,32 +1907,23 @@ void batch_norm_grad(const Tensor& x,
       out_grad.dtype() == phi::DataType::BFLOAT16) {
     out_grad_data = cast<T>(out_grad, phi::DataType::FLOAT32);
   }
-
-  auto x_dims = x_data.dims();
-  const int C = (data_layout_ == DataLayout::kNCHW ? x_dims[1]
-                                                   : x_dims[x_dims.size() - 1]);
-  int nume = 1;
-  for (auto i = 0; i < x_dims.size(); i++) {
-    nume = nume * x_dims[i];
-  }
-
-  const int nhw = nume / C;
-
-  if (x_dims.size() == 2 && data_layout_ == DataLayout::kNCHW) {
+  if (x_data.dims().size() == 2 && data_layout_ == DataLayout::kNCHW) {
     data_layout_ = DataLayout::kNHWC;
   }
-
-  auto run_var = variance_out.get();
-  auto run_mean = mean_out.get();
-
-  Tensor mean_data;
-  Tensor rsqrt_var;
+  auto dtype = x_data.dtype();
 
   if (use_global_stats) {
-    auto eps =
-        full<T>(common::vectorize(run_var.dims()), epsilon, run_var.dtype());
-    mean_data = run_mean;
-    rsqrt_var = rsqrt<T>(run_var + eps);
+    if (has_dynamic_shape(run_var.shape())) {
+      auto eps = backend::full_with_tensor<T>(
+          shape<T>(run_var), epsilon, run_var.dtype());
+      mean_data = run_mean;
+      rsqrt_var = rsqrt<T>(run_var + eps);
+    } else {
+      auto eps =
+          full<T>(common::vectorize(run_var.dims()), epsilon, run_var.dtype());
+      mean_data = run_mean;
+      rsqrt_var = rsqrt<T>(run_var + eps);
+    }
   } else {
     mean_data = saved_mean;
     rsqrt_var = saved_variance;
@@ -1943,11 +1943,6 @@ void batch_norm_grad(const Tensor& x,
   // test mode
   // d_x = d_y * scale * inv_var
 
-  std::vector<int> nchw_to_nhwc_dim = {0, 2, 3, 1};
-  std::vector<int> nhwc_to_nchw_dim = {0, 3, 1, 2};
-  auto reduce_axis = IntArray(std::vector<int64_t>{0, 1, 2});
-  auto dtype = x_data.dtype();
-
   switch (data_layout_) {
     case DataLayout::kNCHW: {
       auto nhwc_x = transpose<T>(x_data, nchw_to_nhwc_dim);
@@ -1973,8 +1968,35 @@ void batch_norm_grad(const Tensor& x,
           if (scale) {
             part1 = scale.get() * part1;
           }
-          auto mean_temp1 = nhwc_out_grad_sum / nhw;
-          auto mean_temp2 = sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          Tensor mean_temp1, mean_temp2;
+          if (has_dynamic_shape_for_x) {
+            Tensor x_dims = shape<T>(x_data);
+            const Tensor C =
+                (data_layout_ == DataLayout::kNCHW
+                     ? get_slice<T>(x_dims, 1)
+                     : get_slice<T>(x_dims, x_data.dims().size() - 1));
+            Tensor nume = full<T>({1}, 1.0, x_dims.dtype());
+            for (auto i = 0; i < x_dims.size(); i++) {
+              nume = nume * get_slice<T>(x_dims, i);
+            }
+            Tensor nhw = nume / C;
+            mean_temp1 = nhwc_out_grad_sum / nhw;
+            mean_temp2 = sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          } else {
+            auto x_dims = x_data.dims();
+            const int C =
+                (data_layout_ == DataLayout::kNCHW ? x_dims[1]
+                                                   : x_dims[x_dims.size() - 1]);
+            int nume = 1;
+            for (auto i = 0; i < x_dims.size(); i++) {
+              nume = nume * x_dims[i];
+            }
+
+            const int nhw = nume / C;
+            mean_temp1 = nhwc_out_grad_sum / nhw;
+            mean_temp2 = sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          }
+
           auto part2 =
               nhwc_out_grad - mean_temp1 - (nhwc_x - mean_data) * mean_temp2;
 
@@ -2015,9 +2037,36 @@ void batch_norm_grad(const Tensor& x,
           if (scale) {
             part1 = scale.get() * part1;
           }
-          auto mean_temp1 = out_grad_data_sum / nhw;
-          auto mean_temp2 =
-              nhwc_sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          Tensor mean_temp1, mean_temp2;
+
+          if (has_dynamic_shape_for_x) {
+            Tensor x_dims = shape<T>(x_data);
+            const Tensor C =
+                (data_layout_ == DataLayout::kNCHW
+                     ? get_slice<T>(x_dims, 1)
+                     : get_slice<T>(x_dims, x_data.dims().size() - 1));
+            Tensor nume = full<T>({1}, 1.0, x_dims.dtype());
+            for (auto i = 0; i < x_dims.size(); i++) {
+              nume = nume * get_slice<T>(x_dims, i);
+            }
+            Tensor nhw = nume / C;
+            mean_temp1 = out_grad_data_sum / nhw;
+            mean_temp2 = nhwc_sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          } else {
+            auto x_dims = x_data.dims();
+            const int C =
+                (data_layout_ == DataLayout::kNCHW ? x_dims[1]
+                                                   : x_dims[x_dims.size() - 1]);
+            int nume = 1;
+            for (auto i = 0; i < x_dims.size(); i++) {
+              nume = nume * x_dims[i];
+            }
+
+            const int nhw = nume / C;
+            mean_temp1 = out_grad_data_sum / nhw;
+            mean_temp2 = nhwc_sum_dout_mul_diff / nhw * rsqrt_var * rsqrt_var;
+          }
+
           auto part2 =
               out_grad_data - mean_temp1 - (x_data - mean_data) * mean_temp2;
 
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -32,6 +32,7 @@
     "pd_op.abs",
     "pd_op.add",
     "pd_op.assign",
+    "pd_op.batch_norm",
     "pd_op.cast",
     "pd_op.concat",
     "pd_op.cos",
diff --git a/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py
@@ -17,6 +17,7 @@
 import numpy as np
 from test_prim_sub_graph_backward_dynamic_shape import (
     TestPrimBaseWithGrad,
+    TestPrimThreeWithGrad,
     TestPrimTwoWithGrad,
     apply_to_static,
 )
@@ -30,6 +31,36 @@ def add_net(x, y):
     return x + y
 
 
+def batch_norm_net1(x, y, z):
+    mean = paddle.ones([40], dtype="float32")
+    var = paddle.zeros([40], dtype='float32')
+    return paddle.nn.functional.batch_norm(x, mean, var, y, z)
+
+
+def batch_norm_net2(x, y, z):
+    mean = paddle.ones([40], dtype="float32")
+    var = paddle.zeros([40], dtype='float32')
+    return paddle.nn.functional.batch_norm(
+        x, mean, var, y, z, use_global_stats=True
+    )
+
+
+def batch_norm_net3(x, y, z):
+    mean = paddle.ones([60], dtype="float32")
+    var = paddle.zeros([60], dtype='float32')
+    return paddle.nn.functional.batch_norm(
+        x, mean, var, y, z, data_format='NHWC'
+    )
+
+
+def batch_norm_net4(x, y, z):
+    mean = paddle.ones([60], dtype="float32")
+    var = paddle.zeros([60], dtype='float32')
+    return paddle.nn.functional.batch_norm(
+        x, mean, var, y, z, use_global_stats=True, data_format='NHWC'
+    )
+
+
 def concat_net1(x):
     y = x + 1
     return paddle.concat([x, y], axis=-1)
@@ -228,6 +259,132 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimBatchNormWithGrad1(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, 40, None, None]
+        self.y_shape = [40]
+        self.init_y_shape = [None]
+        self.z_shape = [40]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad2(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, 40, None, None]
+        self.y_shape = [40]
+        self.init_y_shape = [None]
+        self.z_shape = [40]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad3(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, None, None, 60]
+        self.y_shape = [40]
+        self.init_y_shape = [None]
+        self.z_shape = [40]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad4(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, 40, None, None]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad5(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, None, None, 60]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad6(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, 40, None, None]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net4
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimBatchNormWithGrad7(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50, 60]
+        self.init_x_shape = [None, None, None, 60]
+        self.y_shape = [60]
+        self.init_y_shape = [None]
+        self.z_shape = [60]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.z = np.random.random(self.z_shape).astype(self.dtype)
+        self.net = batch_norm_net4
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimConcatWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py
