[Cleanup][A-13] clean some VarType for test (#61566)

co63oc · web-flow · commit a5bfea86d19e · 2024-02-19T16:07:11.000+08:00
diff --git a/python/paddle/static/amp/decorator.py b/python/paddle/static/amp/decorator.py
@@ -480,7 +480,7 @@ def apply_gradients(self, params_grads):
         found_inf = self._check_finite_and_unscale(params_grads)
         if (
             self._use_dynamic_loss_scaling
-            and self._amp_vartype == core.VarDesc.VarType.FP16
+            and self._amp_vartype == paddle.float16
         ):
             self._add_dynamic_loss_scaling(params_grads, found_inf)
 
@@ -507,7 +507,7 @@ def apply_gradients(self, params_grads):
 
     def _split_grads(self, params_grads):
         grads = [g for _, g in params_grads]
-        fp32_grads = [g for g in grads if g.dtype == core.VarDesc.VarType.FP32]
+        fp32_grads = [g for g in grads if g.dtype == paddle.float32]
         fp16_grads = [g for g in grads if g.dtype == self._amp_vartype]
         assert len(fp32_grads) + len(fp16_grads) == len(
             grads
diff --git a/python/paddle/static/amp/fp16_lists.py b/python/paddle/static/amp/fp16_lists.py
@@ -47,7 +47,7 @@ def check_amp_dtype(dtype):
 
 
 def get_low_precision_vartype(dtype):
-    if isinstance(dtype, core.VarDesc.VarType):
+    if isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
         return dtype
     elif isinstance(dtype, str):
         dtype = dtype.lower()
@@ -71,10 +71,10 @@ def get_low_precision_vartype(dtype):
 def get_low_precision_dtypestr(dtype):
     if isinstance(dtype, str):
         return check_amp_dtype(dtype)
-    elif isinstance(dtype, core.VarDesc.VarType):
-        if dtype == core.VarDesc.VarType.FP16:
+    elif isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
+        if dtype == paddle.float16:
             return "float16"
-        elif dtype == core.VarDesc.VarType.BF16:
+        elif dtype == paddle.bfloat16:
             return "bfloat16"
         else:
             raise ValueError(
diff --git a/python/paddle/static/amp/fp16_utils.py b/python/paddle/static/amp/fp16_utils.py
@@ -185,9 +185,7 @@ def _insert_cast_op(block, op, idx, src_dtype, dest_dtype):
     num_cast_ops = 0
 
     for in_name in op.input_names:
-        if src_dtype == core.VarDesc.VarType.FP32 and _keep_fp32_input(
-            op, in_name
-        ):
+        if src_dtype == paddle.float32 and _keep_fp32_input(op, in_name):
             continue
         for in_var_name in op.input(in_name):
             in_var = block._find_var_recursive(in_var_name)
@@ -210,10 +208,7 @@ def _insert_cast_op(block, op, idx, src_dtype, dest_dtype):
                     # set cast_op device to `all`, can reduce send cast_var.
                     # TODO: need remove this after we unified the dynamic
                     # and static pipeline interface.
-                    if (
-                        src_dtype == core.VarDesc.VarType.FP32
-                        and in_var.stop_gradient
-                    ):
+                    if src_dtype == paddle.float32 and in_var.stop_gradient:
                         prev_op = None
                         if in_var.op is op:
                             prev_op = find_true_prev_op(
@@ -527,7 +522,7 @@ def get_promote_dtype(op, amp_dtype, block):
         if in_name:
             for in_var_name in op.input(in_name):
                 in_var = block._find_var_recursive(in_var_name)
-                if in_var and in_var.dtype == core.VarDesc.VarType.FP32:
+                if in_var and in_var.dtype == paddle.float32:
                     dst_dtype = core.VarDesc.VarType.FP32
                     break
         else:
@@ -915,7 +910,7 @@ def cast_parameters_to_fp16(
             if var_scope.find_var(param.name):
                 param_t = var_scope.find_var(param.name).get_tensor()
                 data = np.array(param_t)
-                if dest_type == core.VarDesc.VarType.BF16:
+                if dest_type == paddle.bfloat16:
                     p_array = _convert_float_to_bfloat16(place, data)
                     param_t.set(p_array, place)
                 else:
@@ -952,7 +947,7 @@ def update_role_var_grad(main_prog, params_grads):
     OPTIMIZE = core.op_proto_and_checker_maker.OpRole.Optimize
     for p, g in params_grads:
         op = g.op
-        if g.dtype == core.VarDesc.VarType.FP32 and op.type == 'cast':
+        if g.dtype == paddle.float32 and op.type == 'cast':
             role = op.attr('op_role')
             if role & int(BACKWARD) and op.has_attr('op_role_var'):
                 op._remove_attr("op_role_var")
diff --git a/python/paddle/static/nn/common.py b/python/paddle/static/nn/common.py
@@ -349,8 +349,8 @@ def instance_norm(
     dtype = helper.input_dtype()
 
     # use fp32 for in parameter
-    if dtype == paddle.framework.core.VarDesc.VarType.FP16:
-        dtype = paddle.framework.core.VarDesc.VarType.FP32
+    if dtype == paddle.float16:
+        dtype = paddle.float32
 
     input_shape = input.shape
     if len(input.shape) < 2 or len(input.shape) > 5:
@@ -2771,8 +2771,8 @@ def batch_norm(
     dtype = helper.input_dtype()
 
     # use fp32 for bn parameter
-    if dtype == core.VarDesc.VarType.FP16 or dtype == core.VarDesc.VarType.BF16:
-        dtype = core.VarDesc.VarType.FP32
+    if dtype == paddle.float16 or dtype == paddle.bfloat16:
+        dtype = paddle.float32
 
     input_shape = input.shape
     if len(input.shape) < 2 or len(input.shape) > 5:
diff --git a/python/paddle/static/quantization/quantization_pass.py b/python/paddle/static/quantization/quantization_pass.py
@@ -516,9 +516,9 @@ def _insert_quant_abs_max_op(
             var_dtype=var_node.dtype(),
         )
         scale_name = self._quantized_scale_name(name)
-        if var_node.dtype() == core.VarDesc.VarType.FP64:
+        if var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -563,9 +563,9 @@ def _insert_quant_range_abs_max_op(
         )
 
         scale_name = self._quantized_scale_name(name)
-        if var_node.dtype() == core.VarDesc.VarType.FP64:
+        if var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -595,9 +595,9 @@ def _insert_quant_range_abs_max_op(
                 shape=[self._window_size],
                 var_dtype=var_node.dtype(),
             )
-            if var_node.dtype() == core.VarDesc.VarType.FP64:
+            if var_node.dtype() == paddle.float64:
                 data_type = 'float64'
-            elif var_node.dtype() == core.VarDesc.VarType.FP32:
+            elif var_node.dtype() == paddle.float32:
                 data_type = 'float32'
             else:
                 data_type = "float16"
@@ -645,9 +645,9 @@ def _insert_quant_moving_average_abs_max_op(
             var_dtype=var_node.dtype(),
         )
         scale_name = self._quantized_scale_name(name)
-        if var_node.dtype() == core.VarDesc.VarType.FP64:
+        if var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -675,9 +675,9 @@ def _insert_quant_moving_average_abs_max_op(
                 var_dtype=var_node.dtype(),
                 shape=[1],
             )
-            if var_node.dtype() == core.VarDesc.VarType.FP64:
+            if var_node.dtype() == paddle.float64:
                 data_type = 'float64'
-            elif var_node.dtype() == core.VarDesc.VarType.FP32:
+            elif var_node.dtype() == paddle.float32:
                 data_type = 'float32'
             else:
                 data_type = "float16"
@@ -753,9 +753,9 @@ def _insert_channel_quant_op(
             var_dtype=var_node.dtype(),
         )
         scale_name = self._quantized_scale_name(name)
-        if var_node.dtype() == core.VarDesc.VarType.FP64:
+        if var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -1291,9 +1291,9 @@ def _insert_post_channel_dequant_op(self, graph, op_node, quant_axis):
             var_dtype=output_var_node.dtype(),
         )
 
-        if output_var_node.dtype() == core.VarDesc.VarType.FP64:
+        if output_var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif output_var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif output_var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -1643,9 +1643,9 @@ def apply(self, graph):
                     ):
                         continue
 
-                    if in_node.dtype() == core.VarDesc.VarType.FP64:
+                    if in_node.dtype() == paddle.float64:
                         data_type = 'float64'
-                    elif in_node.dtype() == core.VarDesc.VarType.FP32:
+                    elif in_node.dtype() == paddle.float32:
                         data_type = 'float32'
                     else:
                         data_type = "float16"
@@ -1997,9 +1997,9 @@ def _inser_quant_dequant_moving_average_abs_max_op(
             var_dtype=var_node.dtype(),
         )
         scale_name = f"{var_node.name()}.quant_dequant@scale"
-        if var_node.dtype() == core.VarDesc.VarType.FP64:
+        if var_node.dtype() == paddle.float64:
             data_type = 'float64'
-        elif var_node.dtype() == core.VarDesc.VarType.FP32:
+        elif var_node.dtype() == paddle.float32:
             data_type = 'float32'
         else:
             data_type = "float16"
@@ -2037,9 +2037,9 @@ def _inser_quant_dequant_moving_average_abs_max_op(
                 var_dtype=var_node.dtype(),
                 shape=[1],
             )
-            if var_node.dtype() == core.VarDesc.VarType.FP64:
+            if var_node.dtype() == paddle.float64:
                 data_type = 'float64'
-            elif var_node.dtype() == core.VarDesc.VarType.FP32:
+            elif var_node.dtype() == paddle.float32:
                 data_type = 'float32'
             else:
                 data_type = "float16"
@@ -2156,9 +2156,9 @@ def insert_quant_op(
             var_dtype=var_node.dtype(),
         )
         if not scale_var_node:
-            if var_node.dtype() == core.VarDesc.VarType.FP64:
+            if var_node.dtype() == paddle.float64:
                 data_type = 'float64'
-            elif var_node.dtype() == core.VarDesc.VarType.FP32:
+            elif var_node.dtype() == paddle.float32:
                 data_type = 'float32'
             else:
                 data_type = "float16"
@@ -2226,9 +2226,9 @@ def insert_quant_op(
                 var_dtype=var_node.dtype(),
                 shape=[1],
             )
-            if var_node.dtype() == core.VarDesc.VarType.FP64:
+            if var_node.dtype() == paddle.float64:
                 data_type = 'float64'
-            elif var_node.dtype() == core.VarDesc.VarType.FP32:
+            elif var_node.dtype() == paddle.float32:
                 data_type = 'float32'
             else:
                 data_type = "float16"
@@ -3419,7 +3419,7 @@ def _insert_quant_dequant_op(self, graph, var_node):
                 )
                 data_type = (
                     'float64'
-                    if var_node.dtype() == core.VarDesc.VarType.FP64
+                    if var_node.dtype() == paddle.float64
                     else 'float32'
                 )
                 _init_var_node(
