[Prim] Fix batch_norm bias_grad loss in cinn (#54751)

cyber-pioneer · web-flow · commit 3e0f0a002b50 · 2023-06-25T10:40:29.000+08:00
* fix batch_norm grad kernel nhwc error

* fix batch_norm bias_grad loss in cinn

* disable cinn

* fix cinn_atol
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -1307,7 +1307,7 @@ void batch_norm_grad(const Tensor& x,
         set_output<T>(scale_grad_data, scale_grad);
       }
       if (bias_grad) {
-        set_output<T>(nhwc_out_grad_sum, bias_grad);
+        set_output<T>(assign<T>(nhwc_out_grad_sum), bias_grad);
       }
       break;
     }
@@ -1343,7 +1343,7 @@ void batch_norm_grad(const Tensor& x,
           set_output<T>(scale_grad_data, scale_grad);
         }
         if (bias_grad) {
-          set_output<T>(out_grad_data_sum, bias_grad);
+          set_output<T>(assign<T>(out_grad_data_sum), bias_grad);
         }
       }
       break;
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -953,7 +953,8 @@ if(WITH_NV_JETSON)
   set_tests_properties(test_conv3d_transpose_op PROPERTIES TIMEOUT 1200)
   set_tests_properties(test_conv3d_op PROPERTIES TIMEOUT 1200)
   set_tests_properties(test_norm_op PROPERTIES TIMEOUT 1200)
-  set_tests_properties(test_batch_norm_op_prim PROPERTIES TIMEOUT 1500)
+  set_tests_properties(test_batch_norm_op_prim_nchw PROPERTIES TIMEOUT 1500)
+  set_tests_properties(test_batch_norm_op_prim_nhwc PROPERTIES TIMEOUT 1500)
   set_tests_properties(test_layer_norm_op PROPERTIES TIMEOUT 1500)
   set_tests_properties(test_pool3d_op PROPERTIES TIMEOUT 1500)
 else()
@@ -962,7 +963,8 @@ else()
   set_tests_properties(test_conv3d_transpose_op PROPERTIES TIMEOUT 120)
   set_tests_properties(test_conv3d_op PROPERTIES TIMEOUT 120)
   set_tests_properties(test_norm_op PROPERTIES TIMEOUT 120)
-  set_tests_properties(test_batch_norm_op_prim PROPERTIES TIMEOUT 250)
+  set_tests_properties(test_batch_norm_op_prim_nchw PROPERTIES TIMEOUT 250)
+  set_tests_properties(test_batch_norm_op_prim_nhwc PROPERTIES TIMEOUT 250)
   set_tests_properties(test_layer_norm_op PROPERTIES TIMEOUT 250)
   set_tests_properties(test_pool3d_op PROPERTIES TIMEOUT 150)
 endif()
@@ -1185,7 +1187,8 @@ set(TEST_CINN_OPS
     test_clip_op
     test_scatter_op
     test_gather_op
-    test_batch_norm_op_prim
+    test_batch_norm_op_prim_nchw
+    test_batch_norm_op_prim_nhwc
     test_layer_norm_op
     test_cast_op
     test_dropout_op
diff --git a/test/legacy_test/test_batch_norm_op_prim_nchw.py b/test/legacy_test/test_batch_norm_op_prim_nchw.py
@@ -108,9 +108,18 @@ def test_check_grad_x(self):
             )
 
     def test_check_grad_scale_bias(self):
-        self.enable_cinn = False
-        self.rev_comp_atol = 1e-3
-        self.rev_comp_rtol = 1e-3
+        if self.data_format == "NCHW":
+            self.enable_cinn = False
+        if self.dtype == "float32":
+            self.rev_comp_atol = 1e-3
+            self.rev_comp_rtol = 1e-3
+            self.cinn_atol = 1e-3
+            self.cinn_rtol = 1e-3
+        elif self.dtype == "float64":
+            self.rev_comp_atol = 1e-12
+            self.rev_comp_rtol = 1e-12
+            self.cinn_atol = 1e-12
+            self.cinn_rtol = 1e-12
         if self.dtype not in ("uint16", "float16"):
             self.check_grad_with_place(
                 core.CPUPlace(),
@@ -252,21 +261,6 @@ def initConfig(self):
         self.use_global_stats = True
 
 
-class TestBatchNormOpNHWCTestMode(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 = [16, 16, 16, 8]
-        self.training = False
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = True
-
-
 class TestBatchNormOpNCHWFp64(TestBatchNormOp):
     def initConfig(self):
         self.fw_comp_atol = 1e-11
@@ -297,21 +291,6 @@ def initConfig(self):
         self.use_global_stats = None
 
 
-class TestBatchNormOpNHWCTestModeFp64(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-15
-        self.fw_comp_rtol = 1e-15
-        self.rev_comp_atol = 1e-15
-        self.rev_comp_rtol = 1e-15
-        self.dtype = "float64"
-        self.shape = [16, 16, 16, 8]
-        self.training = False
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
 class TestBatchNormOpNCHWFp16(TestBatchNormOp):
     def initConfig(self):
         self.fw_comp_atol = 1e-3
@@ -342,21 +321,6 @@ def initConfig(self):
         self.use_global_stats = None
 
 
-class TestBatchNormOpNHWCTestModeFp16(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-3
-        self.fw_comp_rtol = 1e-3
-        self.rev_comp_atol = 1e-3
-        self.rev_comp_rtol = 1e-3
-        self.dtype = "float16"
-        self.shape = [16, 16, 16, 8]
-        self.training = False
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_bfloat16_supported(core.CUDAPlace(0)),
@@ -401,95 +365,6 @@ def initConfig(self):
         self.use_global_stats = None
 
 
-@unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
-    "core is not compiled with CUDA or not support the bfloat16",
-)
-class TestBatchNormOpNHWCTestModebf16(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-3
-        self.fw_comp_rtol = 1e-3
-        self.rev_comp_atol = 1e-3
-        self.rev_comp_rtol = 1e-3
-        self.cinn_atol = 1e-3
-        self.cinn_rtol = 1e-3
-        self.dtype = "uint16"
-        self.shape = [16, 16, 16, 8]
-        self.training = False
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-class TestBatchNormOpNHWC(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 = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-class TestBatchNormOpNHWCFp64(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-11
-        self.fw_comp_rtol = 1e-11
-        self.rev_comp_atol = 1e-11
-        self.rev_comp_rtol = 1e-11
-        self.dtype = "float64"
-        self.shape = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-class TestBatchNormOpNHWCFp16(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-3
-        self.fw_comp_rtol = 1e-3
-        self.rev_comp_atol = 1e-3
-        self.rev_comp_rtol = 1e-3
-        self.dtype = "float16"
-        self.shape = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-@unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
-    "core is not compiled with CUDA or not support the bfloat16",
-)
-class TestBatchNormOpNHWCbf16(TestBatchNormOp):
-    def initConfig(self):
-        self.fw_comp_atol = 1e-3
-        self.fw_comp_rtol = 1e-3
-        self.rev_comp_atol = 1e-3
-        self.rev_comp_rtol = 1e-3
-        self.cinn_atol = 1e-3
-        self.cinn_rtol = 1e-3
-        self.dtype = "uint16"
-        self.shape = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
 class TestBatchNormOpNCHWShape2(TestBatchNormOp):
     def initConfig(self):
         self.fw_comp_atol = 1e-5
@@ -535,51 +410,6 @@ def initConfig(self):
         self.use_global_stats = None
 
 
-class TestBatchNormOpNHWCShape2(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, 16, 32]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-class TestBatchNormOpNHWCMomentum2(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 = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.9
-        self.epsilon = 1e-05
-        self.data_format = "NHWC"
-        self.use_global_stats = None
-
-
-class TestBatchNormOpNHWCEps2(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 = [16, 16, 16, 8]
-        self.training = True
-        self.momentum = 0.1
-        self.epsilon = 1e-06
-        self.data_format = "NHWC"
-        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

Original file line number	Diff line number	Diff line change
`@@ -1307,7 +1307,7 @@ void batch_norm_grad(const Tensor& x,`
`1307`	`1307`	`set_output<T>(scale_grad_data, scale_grad);`
`1308`	`1308`	`}`
`1309`	`1309`	`if (bias_grad) {`
`1310`		`- set_output<T>(nhwc_out_grad_sum, bias_grad);`
	`1310`	`+ set_output<T>(assign<T>(nhwc_out_grad_sum), bias_grad);`
`1311`	`1311`	`}`
`1312`	`1312`	`break;`
`1313`	`1313`	`}`
`@@ -1343,7 +1343,7 @@ void batch_norm_grad(const Tensor& x,`
`1343`	`1343`	`set_output<T>(scale_grad_data, scale_grad);`
`1344`	`1344`	`}`
`1345`	`1345`	`if (bias_grad) {`
`1346`		`- set_output<T>(out_grad_data_sum, bias_grad);`
	`1346`	`+ set_output<T>(assign<T>(out_grad_data_sum), bias_grad);`
`1347`	`1347`	`}`
`1348`	`1348`	`}`
`1349`	`1349`	`break;`