add support for batchnorm-1d (#69054)

cubehan3 · web-flow · commit 1aa31f86bbc8 · 2024-10-31T16:34:11.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1816,7 +1816,9 @@ void batch_norm_grad(const Tensor& x,
   use_global_stats = is_test || use_global_stats;
 
   bool has_dynamic_shape_for_x = has_dynamic_shape(x.shape());
-
+  bool has_dynamic_shape_for_out_grad = has_dynamic_shape(out_grad.shape());
+  bool dim_three = x.dims().size() == 3;
+  bool dim_two = x.dims().size() == 2;
   DataLayout data_layout_ = common::StringToDataLayout(data_layout);
 
   Tensor x_data = ConverToMT<T>(x);
@@ -1829,9 +1831,71 @@ void batch_norm_grad(const Tensor& x,
   std::vector<int> nhwc_to_nchw_dim = {0, 3, 1, 2};
   auto reduce_axis = IntArray(std::vector<int64_t>{0, 1, 2});
 
-  if (x_data.dims().size() == 2 && data_layout_ == DataLayout::kNCHW) {
-    data_layout_ = DataLayout::kNHWC;
+  if (dim_two) {
+    if (data_layout_ == DataLayout::kNCHW) {
+      data_layout_ = DataLayout::kNHWC;
+    }
+
+    if (has_dynamic_shape_for_x) {
+      x_data = backend::reshape<T>(
+          x_data, get_unsqueeze_dims<T>(shape<T>(x_data), {1, 2}));
+    } else {
+      x_data = reshape<T>(x_data, get_unsqueeze_dims(x_data, {1, 2}));
+    }
+
+    if (has_dynamic_shape_for_out_grad) {
+      out_grad_data = backend::reshape<T>(
+          out_grad_data,
+          get_unsqueeze_dims<T>(shape<T>(out_grad_data), {1, 2}));
+    } else {
+      out_grad_data =
+          reshape<T>(out_grad_data, get_unsqueeze_dims(out_grad_data, {1, 2}));
+    }
+  } else if (dim_three) {
+    // Add an additional axis to accommodate NCHW or NHWC formats.
+    switch (data_layout_) {
+      case DataLayout::kNCHW: {
+        if (has_dynamic_shape_for_x) {
+          x_data = backend::reshape<T>(
+              x_data, get_unsqueeze_dims<T>(shape<T>(x_data), {3}));
+        } else {
+          x_data = reshape<T>(x_data, get_unsqueeze_dims(x_data, {3}));
+        }
+
+        if (has_dynamic_shape_for_out_grad) {
+          out_grad_data = backend::reshape<T>(
+              out_grad_data,
+              get_unsqueeze_dims<T>(shape<T>(out_grad_data), {3}));
+        } else {
+          out_grad_data =
+              reshape<T>(out_grad_data, get_unsqueeze_dims(out_grad_data, {3}));
+        }
+        break;
+      }
+      case DataLayout::kNHWC: {
+        if (has_dynamic_shape_for_x) {
+          x_data = backend::reshape<T>(
+              x_data, get_unsqueeze_dims<T>(shape<T>(x_data), {2}));
+        } else {
+          x_data = reshape<T>(x_data, get_unsqueeze_dims(x_data, {2}));
+        }
+
+        if (has_dynamic_shape_for_out_grad) {
+          out_grad_data = backend::reshape<T>(
+              out_grad_data,
+              get_unsqueeze_dims<T>(shape<T>(out_grad_data), {2}));
+        } else {
+          out_grad_data =
+              reshape<T>(out_grad_data, get_unsqueeze_dims(out_grad_data, {2}));
+        }
+        break;
+      }
+      default:
+        PADDLE_THROW(common::errors::InvalidArgument(
+            "Unknown storage order: %s", data_layout));
+    }
   }
+
   auto dtype = x_data.dtype();
 
   if (use_global_stats) {
@@ -1883,6 +1947,13 @@ void batch_norm_grad(const Tensor& x,
           auto nchw_x_grad = transpose<T>(nhwc_x_grad, nhwc_to_nchw_dim);
           nchw_x_grad = ConverToOrig<T>(nchw_x_grad, x.dtype());
 
+          if (dim_three) {
+            if (has_dynamic_shape_for_x) {
+              nchw_x_grad = backend::reshape<T>(nchw_x_grad, shape<T>(x));
+            } else {
+              nchw_x_grad = reshape<T>(nchw_x_grad, x.shape());
+            }
+          }
           set_output<T>(nchw_x_grad, x_grad);
         } else {
           auto part1 = rsqrt_var;
@@ -1924,6 +1995,14 @@ void batch_norm_grad(const Tensor& x,
           auto x_grad_data = part1 * part2;
           auto nchw_x_grad = transpose<T>(x_grad_data, nhwc_to_nchw_dim);
           nchw_x_grad = ConverToOrig<T>(nchw_x_grad, x.dtype());
+
+          if (dim_three) {
+            if (has_dynamic_shape_for_x) {
+              nchw_x_grad = backend::reshape<T>(nchw_x_grad, shape<T>(x));
+            } else {
+              nchw_x_grad = reshape<T>(nchw_x_grad, x.shape());
+            }
+          }
           set_output<T>(nchw_x_grad, x_grad);
         }
       }
@@ -1947,6 +2026,13 @@ void batch_norm_grad(const Tensor& x,
           if (scale) {
             x_grad_data = scale.get() * x_grad_data;
           }
+          if (dim_two || dim_three) {
+            if (has_dynamic_shape_for_x) {
+              x_grad_data = backend::reshape<T>(x_grad_data, shape<T>(x));
+            } else {
+              x_grad_data = reshape<T>(x_grad_data, x.shape());
+            }
+          }
           x_grad_data = ConverToOrig<T>(x_grad_data, x.dtype());
           set_output<T>(x_grad_data, x_grad);
         } else {
@@ -1989,7 +2075,13 @@ void batch_norm_grad(const Tensor& x,
 
           auto x_grad_data = part1 * part2;
           x_grad_data = ConverToOrig<T>(x_grad_data, x.dtype());
-
+          if (dim_two || dim_three) {
+            if (has_dynamic_shape_for_x) {
+              x_grad_data = backend::reshape<T>(x_grad_data, shape<T>(x));
+            } else {
+              x_grad_data = reshape<T>(x_grad_data, x.shape());
+            }
+          }
           set_output<T>(x_grad_data, x_grad);
         }
         if (scale_grad) {
diff --git a/test/legacy_test/test_batch_norm_op_prim_nchw.py b/test/legacy_test/test_batch_norm_op_prim_nchw.py
@@ -432,6 +432,36 @@ def initConfig(self):
         self.use_global_stats = None
 
 
+class TestBatchNormOpNCHWShape3(TestBatchNormOp):
+    def initConfig(self):
+        self.fw_comp_atol = 1e-5
+        self.fw_comp_rtol = 1e-5
+        self.rev_comp_atol = 1e-5
+        self.rev_comp_rtol = 1e-5
+        self.dtype = "float32"
+        self.shape = [4, 8, 32]
+        self.training = True
+        self.momentum = 0.1
+        self.epsilon = 1e-05
+        self.data_format = "NCHW"
+        self.use_global_stats = None
+
+
+class TestBatchNormOpNCHWShape4(TestBatchNormOp):
+    def initConfig(self):
+        self.fw_comp_atol = 1e-5
+        self.fw_comp_rtol = 1e-5
+        self.rev_comp_atol = 1e-5
+        self.rev_comp_rtol = 1e-5
+        self.dtype = "float32"
+        self.shape = [4, 256]
+        self.training = True
+        self.momentum = 0.1
+        self.epsilon = 1e-05
+        self.data_format = "NCHW"
+        self.use_global_stats = None
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
diff --git a/test/legacy_test/test_batch_norm_op_prim_nhwc.py b/test/legacy_test/test_batch_norm_op_prim_nhwc.py
@@ -221,6 +221,36 @@ def initConfig(self):
         self.use_global_stats = None
 
 
+class TestBatchNormOpNHWCShape3(TestBatchNormOp):
+    def initConfig(self):
+        self.fw_comp_atol = 1e-5
+        self.fw_comp_rtol = 1e-5
+        self.rev_comp_atol = 1e-5
+        self.rev_comp_rtol = 1e-5
+        self.dtype = "float32"
+        self.shape = [4, 128, 32]
+        self.training = True
+        self.momentum = 0.1
+        self.epsilon = 1e-05
+        self.data_format = "NHWC"
+        self.use_global_stats = None
+
+
+class TestBatchNormOpNHWCShape4(TestBatchNormOp):
+    def initConfig(self):
+        self.fw_comp_atol = 1e-5
+        self.fw_comp_rtol = 1e-5
+        self.rev_comp_atol = 1e-5
+        self.rev_comp_rtol = 1e-5
+        self.dtype = "float32"
+        self.shape = [4, 256]
+        self.training = True
+        self.momentum = 0.1
+        self.epsilon = 1e-05
+        self.data_format = "NHWC"
+        self.use_global_stats = None
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
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
@@ -410,6 +410,82 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimBatchNormWithGrad8(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.batch_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50]
+        self.init_x_shape = [None, 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-5
+
+
+class TestPrimBatchNormWithGrad9(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.batch_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 40]
+        self.init_x_shape = [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-5
+
+
+class TestPrimBatchNormWithGrad10(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.batch_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 60]
+        self.init_x_shape = [None, 40, 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-5
+
+
+class TestPrimBatchNormWithGrad11(TestPrimThreeWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.batch_norm_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 60]
+        self.init_x_shape = [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-5
+
+
 class TestPrimConcatWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)
