[NPU][MLU] Fix test_batch_nrom & test_LeNet_MNIST & test_custom_pass (#1626)

Author: LittleHeroZZZX

backends/mlu/tests/test_LeNet_MNIST.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import os
-import shutil
 import time
 import argparse
 import datetime
+import tempfile
 import numpy as np
 
 import paddle
@@ -67,13 +67,9 @@ def test(epoch_id, test_loader, model, cost):
     )
 
 
-def infer(model_dir):
-    # model file
-    params_file = os.path.join(model_dir, "model.pdiparams")
-    model_file = os.path.join(model_dir, "model.pdmodel")
-
+def infer(model_dir, prefix):
     # create config
-    config = paddle_infer.Config(model_file, params_file)
+    config = paddle_infer.Config(model_dir, prefix)
     config.enable_custom_device("mlu")
 
     # create predictor
@@ -203,11 +199,13 @@ def main(args):
         model,
         input_spec=[paddle.static.InputSpec(shape=[None, 1, 28, 28], dtype="float32")],
     )
-    paddle.jit.save(model, "output/model")
 
-    # infernece and clear
-    infer("output")
-    shutil.rmtree("output")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        prefix = "test_LeNet_MNIST"
+        paddle.jit.save(model, os.path.join(temp_dir, prefix))
+
+        # infernece and clear
+        infer(temp_dir, prefix)
 
 
 class AverageMeter(object):

backends/mlu/tests/unittests/test_batch_norm_op_mlu.py

@@ -478,81 +478,82 @@ def test_with_place(place, data_layout, shape):
             ]
             ground_truth = {name: var_dict[name] for name in var_names}
 
-            program = base.Program()
-            with base.program_guard(program):
-                block = program.global_block()
-                for name in ground_truth:
-                    block.create_var(
-                        name=name, dtype="float32", shape=ground_truth[name].shape
+            with paddle.pir_utils.OldIrGuard():
+                program = base.Program()
+                with base.program_guard(program):
+                    block = program.global_block()
+                    for name in ground_truth:
+                        block.create_var(
+                            name=name, dtype="float32", shape=ground_truth[name].shape
+                        )
+                    inputs = {
+                        "X": block.var("x"),
+                        "Scale": block.var("scale"),
+                        "Bias": block.var("bias"),
+                        "Mean": block.var("mean"),
+                        "Variance": block.var("variance"),
+                    }
+                    attrs = {
+                        "epsilon": epsilon,
+                        "is_test": False,
+                        "data_layout": data_layout,
+                        "use_mkldnn": False,
+                        "fuse_with_relu": self.fuse_with_relu,
+                        "use_global_stats": self.use_global_stats,
+                    }
+                    if self.use_momentum_variable:
+                        inputs["MomentumTensor"] = block.var("momentum_var")
+                    else:
+                        attrs["momentum"] = momentum
+
+                    outputs = {
+                        "Y": block.var("y"),
+                        "MeanOut": block.var("mean"),  # share memory
+                        "VarianceOut": block.var("variance"),  # share memory
+                        "SavedMean": block.var("saved_mean"),
+                        "SavedVariance": block.var("saved_variance"),
+                    }
+                    block.create_var(name="reserve_space", dtype="float32")
+                    outputs["ReserveSpace"] = block.var("reserve_space")
+                    bn_op = block.append_op(
+                        type="batch_norm", inputs=inputs, outputs=outputs, attrs=attrs
+                    )
+                    block.create_var(name="y@GRAD", dtype="float32", shape=y.shape)
+
+                    # generate backward op_desc
+                    grad_op_desc_list, op_grad_to_var = core.get_grad_op_desc(
+                        bn_op.desc, self.no_grad_set, []
+                    )
+                    grad_op_desc = grad_op_desc_list[0]
+                    new_op_desc = block.desc.append_op()
+                    new_op_desc.copy_from(grad_op_desc)
+                    for var_name in grad_op_desc.output_arg_names():
+                        block.desc.var(var_name.encode("ascii"))
+                    grad_op_desc.infer_var_type(block.desc)
+                    grad_op_desc.infer_shape(block.desc)
+                    for arg in grad_op_desc.output_arg_names():
+                        grad_var = block.desc.find_var(arg.encode("ascii"))
+                        grad_var.set_dtype(core.VarDesc.VarType.FP32)
+
+                    program._sync_with_cpp()
+
+                    exe = base.Executor(place)
+                    out = exe.run(
+                        program,
+                        feed={
+                            name: var_dict[name]
+                            for name in [
+                                "x",
+                                "scale",
+                                "bias",
+                                "mean",
+                                "variance",
+                                "y@GRAD",
+                                "momentum_var",
+                            ]
+                        },
+                        fetch_list=self.fetch_list,
                     )
-                inputs = {
-                    "X": block.var("x"),
-                    "Scale": block.var("scale"),
-                    "Bias": block.var("bias"),
-                    "Mean": block.var("mean"),
-                    "Variance": block.var("variance"),
-                }
-                attrs = {
-                    "epsilon": epsilon,
-                    "is_test": False,
-                    "data_layout": data_layout,
-                    "use_mkldnn": False,
-                    "fuse_with_relu": self.fuse_with_relu,
-                    "use_global_stats": self.use_global_stats,
-                }
-                if self.use_momentum_variable:
-                    inputs["MomentumTensor"] = block.var("momentum_var")
-                else:
-                    attrs["momentum"] = momentum
-
-                outputs = {
-                    "Y": block.var("y"),
-                    "MeanOut": block.var("mean"),  # share memory
-                    "VarianceOut": block.var("variance"),  # share memory
-                    "SavedMean": block.var("saved_mean"),
-                    "SavedVariance": block.var("saved_variance"),
-                }
-                block.create_var(name="reserve_space", dtype="float32")
-                outputs["ReserveSpace"] = block.var("reserve_space")
-                bn_op = block.append_op(
-                    type="batch_norm", inputs=inputs, outputs=outputs, attrs=attrs
-                )
-                block.create_var(name="y@GRAD", dtype="float32", shape=y.shape)
-
-                # generate backward op_desc
-                grad_op_desc_list, op_grad_to_var = core.get_grad_op_desc(
-                    bn_op.desc, self.no_grad_set, []
-                )
-                grad_op_desc = grad_op_desc_list[0]
-                new_op_desc = block.desc.append_op()
-                new_op_desc.copy_from(grad_op_desc)
-                for var_name in grad_op_desc.output_arg_names():
-                    block.desc.var(var_name.encode("ascii"))
-                grad_op_desc.infer_var_type(block.desc)
-                grad_op_desc.infer_shape(block.desc)
-                for arg in grad_op_desc.output_arg_names():
-                    grad_var = block.desc.find_var(arg.encode("ascii"))
-                    grad_var.set_dtype(core.VarDesc.VarType.FP32)
-
-                program._sync_with_cpp()
-
-                exe = base.Executor(place)
-                out = exe.run(
-                    program,
-                    feed={
-                        name: var_dict[name]
-                        for name in [
-                            "x",
-                            "scale",
-                            "bias",
-                            "mean",
-                            "variance",
-                            "y@GRAD",
-                            "momentum_var",
-                        ]
-                    },
-                    fetch_list=self.fetch_list,
-                )
 
             for id, name in enumerate(self.fetch_list):
                 if name == "variance":

backends/mlu/tests/unittests/test_custom_pass_mlu.py

@@ -14,6 +14,7 @@
 
 from __future__ import print_function, division
 
+import tempfile
 import os
 import numpy as np
 import unittest
@@ -33,33 +34,54 @@ def replace(x, y, z):
     return pattern, replace
 
 
-@paddle.jit.to_static(
-    input_spec=[
-        paddle.static.InputSpec([None, 32], "float32", "x"),
-        paddle.static.InputSpec([None, 32], "float32", "y"),
-        paddle.static.InputSpec([None, 32], "float32", "z"),
-    ]
-)
-def func(x, y, z):
-    return x + y + z
+class TestNet(paddle.nn.Layer):
+    def __init__(self):
+        super(TestNet, self).__init__()
+        if paddle.framework.use_pir_api():
+            # In PIR, .pdiparams file is necessary when loading model.
+            # At least one parameter is needed to generate .pdiparams file in paddle.jit.save
+            self.alpha = paddle.create_parameter(shape=[1], dtype="float32")
 
-
-MODLE_FILE = "./saved_model"
+    @paddle.jit.to_static(
+        full_graph=True,
+        input_spec=[
+            paddle.static.InputSpec([None, 32], "float32", "x"),
+            paddle.static.InputSpec([None, 32], "float32", "y"),
+            paddle.static.InputSpec([None, 32], "float32", "z"),
+        ],
+    )
+    def forward(self, x, y, z):
+        return x + y + z
 
 
 class TestCustomPass(unittest.TestCase):
     def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.prefix = "test_custom_pass"
+
         for lib in os.listdir(os.getenv("CUSTOM_DEVICE_ROOT")):
             if lib.endswith(".so"):
                 paddle.utils.cpp_extension.extension_utils.load_op_meta_info_and_register_op(
                     lib
                 )
-        paddle.jit.save(func, MODLE_FILE)
+
+        paddle.disable_static()
+        net = TestNet()
+        paddle.enable_static()
+        paddle.jit.save(net, os.path.join(self.temp_dir.name, self.prefix))
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
 
     def test_my_add_n(self):
-        config = paddle.inference.Config()
-        config.set_prog_file(MODLE_FILE + ".pdmodel")
-        config.enable_memory_optim()
+        if paddle.framework.use_pir_api():
+            config = paddle.inference.Config(self.temp_dir.name, self.prefix)
+        else:
+            config = paddle.inference.Config()
+            config.set_prog_file(
+                os.path.join(self.temp_dir.name, self.prefix + ".pdmodel")
+            )
+            config.enable_memory_optim()
         config.enable_custom_device("mlu")
         pass_builder = config.pass_builder()
         pass_builder.append_pass("generate_add_n")

backends/npu/tests/test_LeNet_MNIST.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import os
-import shutil
 import time
 import argparse
 import datetime
+import tempfile
 import numpy as np
 
 import paddle
@@ -67,13 +67,9 @@ def test(epoch_id, test_loader, model, cost):
     )
 
 
-def infer(model_dir):
-    # model file
-    params_file = os.path.join(model_dir, "model.pdiparams")
-    model_file = os.path.join(model_dir, "model.pdmodel")
-
+def infer(model_dir, prefix):
     # create config
-    config = paddle_infer.Config(model_file, params_file)
+    config = paddle_infer.Config(model_dir, prefix)
     config.enable_custom_device("npu")
 
     # create predictor
@@ -203,11 +199,12 @@ def main(args):
         model,
         input_spec=[paddle.static.InputSpec(shape=[None, 1, 28, 28], dtype="float32")],
     )
-    paddle.jit.save(model, "output/model")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        prefix = "test_LeNet_MNIST"
+        paddle.jit.save(model, os.path.join(temp_dir, prefix))
 
-    # infernece and clear
-    infer("output")
-    shutil.rmtree("output")
+        # infernece and clear
+        infer(temp_dir, prefix)
 
 
 class AverageMeter(object):