fixed bugs in batchnorm backward decomp (#69688)

cubehan3 · web-flow · commit b8871d79b6d2 · 2024-11-26T10:05:10.000+08:00
diff --git a/paddle/fluid/primitive/decomp_utils/decomp_utils.h b/paddle/fluid/primitive/decomp_utils/decomp_utils.h
@@ -447,7 +447,7 @@ class BatchNormDecompHelper {
     } else {
       auto x_shape = shape<T>(x);
       auto nhw = get_slice<T>(x_shape, 0);
-      for (int64_t i = 0; i < x_rank_; ++i) {
+      for (int64_t i = 1; i < x_rank_; ++i) {
         if (i == channel_axis_) {
           continue;
         }
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
@@ -69,6 +69,14 @@ def batch_norm_net4(x, y, z):
     )
 
 
+def batch_norm_net5(x, y, z):
+    var = paddle.ones([40], dtype="float32")
+    mean = paddle.zeros([40], dtype='float32')
+    return paddle.nn.functional.batch_norm(
+        x, mean, var, y, z, use_global_stats=False, training=True
+    )
+
+
 def ceil_net(x):
     return paddle.ceil(x)
 
@@ -524,6 +532,25 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimBatchNormWithGrad12(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_net5
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 class TestPrimCeilWithGrad(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2024)
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -1143,6 +1143,27 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimBatchNormNC(TestPrimThree):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [30, 40]
+        self.shape_y = [40]
+        self.shape_z = [40]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.dtype_z = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None]
+        self.init_z_shape = [None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.y = np.random.random(self.shape_y).astype(self.dtype_y)
+        self.z = np.random.random(self.shape_z).astype(self.dtype_z)
+        self.net = batch_norm_net2
+        self.necessary_ops = "pd_op.batch_norm_"
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 class TestPrimLogLoss1(TestPrimTwo):
     def setUp(self):
         np.random.seed(2023)

Original file line number	Diff line number	Diff line change
`@@ -447,7 +447,7 @@ class BatchNormDecompHelper {`
`447`	`447`	`} else {`
`448`	`448`	`auto x_shape = shape<T>(x);`
`449`	`449`	`auto nhw = get_slice<T>(x_shape, 0);`
`450`		`- for (int64_t i = 0; i < x_rank_; ++i) {`
	`450`	`+ for (int64_t i = 1; i < x_rank_; ++i) {`
`451`	`451`	`if (i == channel_axis_) {`
`452`	`452`	`continue;`
`453`	`453`	`}`