update test

python/paddle/fluid/tests/unittests/test_batch_norm_op.py

@@ -187,11 +187,13 @@ def set_output_grad(scope, outputs, place, feed_dict=None):
 
 
 class TestBatchNormOpInference(OpTest):
+    def setUp(self):
+        self.dtype = np.float32
+
     def __assert_close(self, tensor, np_array, msg, atol=1e-4):
         self.assertTrue(np.allclose(np.array(tensor), np_array, atol=atol), msg)
 
-    def test_inference(self):
-        def test_with_place(place, data_layout, dtype, shape):
+    def check_with_place(place, data_layout, dtype, shape):
         epsilon = 0.00001
         if len(shape) == 2:
             x_shape = shape
@@ -213,33 +215,32 @@ class TestBatchNormOpInference(OpTest):
         mean = np.zeros(scale_shape).astype(dtype)
         variance = np.ones(scale_shape).astype(dtype)
 
-            y_out = _reference_testing(x_val, scale_val, bias_val, mean,
-                                       variance, epsilon,
-                                       data_layout).astype(dtype)
+        y_out = _reference_testing(x_val, scale_val, bias_val, mean, variance,
+                                   epsilon, data_layout).astype(dtype)
 
         scope = core.Scope()
 
         # create input
-            x_tensor = create_or_get_tensor(
-                scope, "x_val", OpTest.np_dtype_to_fluid_dtype(x_val), place)
+        x_tensor = create_or_get_tensor(scope, "x_val",
+                                        OpTest.np_dtype_to_fluid_dtype(x_val),
+                                        place)
         scale_tensor = create_or_get_tensor(
             scope, "scale_val",
             OpTest.np_dtype_to_fluid_dtype(scale_val), place)
         bias_tensor = create_or_get_tensor(
-                scope, "bias_val",
-                OpTest.np_dtype_to_fluid_dtype(bias_val), place)
-            mean_tensor = create_or_get_tensor(
-                scope, "mean", OpTest.np_dtype_to_fluid_dtype(mean), place)
+            scope, "bias_val", OpTest.np_dtype_to_fluid_dtype(bias_val), place)
+        mean_tensor = create_or_get_tensor(scope, "mean",
+                                           OpTest.np_dtype_to_fluid_dtype(mean),
+                                           place)
         variance_tensor = create_or_get_tensor(
-                scope, "variance",
-                OpTest.np_dtype_to_fluid_dtype(variance), place)
+            scope, "variance", OpTest.np_dtype_to_fluid_dtype(variance), place)
 
         # create output
         y_tensor = create_or_get_tensor(scope, "y_out", None, place)
         saved_mean_tensor = create_or_get_tensor(scope, "saved_mean", None,
                                                  place)
-            saved_variance_tensor = create_or_get_tensor(
-                scope, "saved_variance", None, place)
+        saved_variance_tensor = create_or_get_tensor(scope, "saved_variance",
+                                                     None, place)
         mean_out_tensor = mean_tensor
         variance_out_tensor = variance_tensor
 
@@ -265,21 +266,37 @@ class TestBatchNormOpInference(OpTest):
         batch_norm_op.run(scope, place)
 
         # check inference result
-            self.__assert_close(
-                y_tensor, y_out, "inference output are different at " +
-                str(place) + ", " + data_layout + ", " + str(np.dtype(dtype)))
+        self.__assert_close(y_tensor, y_out,
+                            "inference output are different at " + str(place) +
+                            ", " + data_layout + ", " + str(np.dtype(dtype)))
 
+    def test_check_output(self):
         places = [core.CPUPlace()]
+        if core.is_compiled_with_cuda() and core.op_support_gpu("batch_norm"):
+            places.append(core.CUDAPlace(0))
+
+        for place in places:
+            for data_format in ["NCHW", "NHWC"]:
+                check_with_place(place, data_format, self.dtype, [2, 3, 4, 5])
+                check_with_place(place, data_format, self.dtype, [2, 3])
+
+
+class TestFP16BatchNormOpInference(TestBatchNormOpInference):
+    def setUp(self):
+        self.dtype = np.float16
+
+    def test_check_output(self):
+        places = []
         if core.is_compiled_with_cuda() and core.op_support_gpu("batch_norm"):
             place = core.CUDAPlace(0)
-            if self.dtype != np.float16 or core.is_float16_supported(place):
+            if core.is_float16_supported(place):
                 places.append(place)
 
         for place in places:
             for data_format in ["NCHW", "NHWC"]:
-                for dtype in [np.float32, np.float16]:
-                    test_with_place(place, data_format, dtype, [2, 3, 4, 5])
-                    test_with_place(place, data_format, dtype, [2, 3])
+                check_output_with_place(place, data_format, self.dtype,
+                                        [2, 3, 4, 5])
+                check_output_with_place(place, data_format, self.dtype, [2, 3])
 
 
 class TestBatchNormOpTraining(OpTest):