Fix test_sot_export unittest & add create inputs interface (#61100)

feifei-111 · web-flow · commit 05ccbe84d49a · 2024-01-25T19:50:16.000+08:00
* update

* update

* create input

* update
diff --git a/python/paddle/jit/sot/symbolic/export.py b/python/paddle/jit/sot/symbolic/export.py
@@ -130,6 +130,8 @@ def gen_py_codes(self):
         self.new_root("\n")
         self.create_layer()
         self.new_root("\n")
+        self.create_inputs()
+        self.new_root("\n")
         self.create_test()
         self.new_root("\n")
         self.create_tail()
@@ -201,13 +203,14 @@ def create_layer(self):
             )
         )
 
-    def create_test(self):
-        test_class = self.new_root("class TestLayer(unittest.TestCase):")
+    def create_inputs(self):
+        create_paddle_inputs = self.new_root("def create_paddle_inputs():")
+        self.new_root("\n")
+        craete_numpy_inputs = self.new_root("def create_numpy_inputs():")
+
+        paddle_inputs = ["inputs = ("]
+        numpy_inputs = ["inputs = ("]
 
-        setup = test_class.add_sub("def setUp(self):")
-        test_inputs = [
-            "self.inputs = (",
-        ]
         for inp in self.SIR.inputs:
             if inp in self.SIR.non_param_symbol:
                 meta = self.SIR.symbol_meta_map[inp.name]
@@ -218,23 +221,45 @@ def create_test(self):
                     paddle.int32,
                     paddle.int64,
                 ):
-                    test_inputs.append(
+                    paddle_inputs.append(
                         f"    paddle.randint(low=0, high=10, shape={shape_str}, dtype={meta.dtype}),"
                     )
+                    numpy_inputs.append(
+                        "    np.random.randint(low=0, high=10, size={}, dtype='{}'),".format(
+                            shape_str, str(meta.dtype).replace('paddle.', '')
+                        )
+                    )
                 else:
-                    test_inputs.append(
+                    paddle_inputs.append(
                         f"    paddle.rand(shape={shape_str}, dtype={meta.dtype}),"
                     )
-        test_inputs.append(")")
-        setup.add_sub(*test_inputs)
+                    numpy_inputs.append(
+                        "    np.random.random(size={}).astype('{}'),".format(
+                            shape_str, str(meta.dtype).replace('paddle.', '')
+                        )
+                    )
+
+        paddle_inputs.append(")")
+        paddle_inputs.append("return inputs")
+        numpy_inputs.append(")")
+        numpy_inputs.append("return inputs")
+
+        create_paddle_inputs.add_sub(*paddle_inputs)
+        craete_numpy_inputs.add_sub(*numpy_inputs)
+
+    def create_test(self):
+        test_class = self.new_root("class TestLayer(unittest.TestCase):")
+
+        setup = test_class.add_sub("def setUp(self):")
+        setup.add_sub("self.inputs = create_paddle_inputs()")
         setup.add_sub("self.net = LayerCase()")
 
         train = test_class.add_sub(
             "def train(self, net, to_static, with_prim=False, with_cinn=False):"
         )
         train.add_sub(
-            "paddle.set_flags({'FLAGS_prim_all': with_prim})",
             "if to_static:",
+            "    paddle.set_flags({'FLAGS_prim_all': with_prim})",
             "    if with_cinn:",
             "        build_strategy = paddle.static.BuildStrategy()",
             "        build_strategy.build_cinn_pass = True",
diff --git a/test/sot/test_sot_export.py b/test/sot/test_sot_export.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import tempfile
 import unittest
 
 import paddle
@@ -37,12 +38,15 @@ def forward(self, x):
 
 class TestSotExport(unittest.TestCase):
     @min_graph_size_guard(0)
-    @with_export_guard("/tmp")
     def test_basic(self):
+        temp_dir = tempfile.TemporaryDirectory()
+        temp_dir_name = temp_dir.name
         net = Net()
         x = paddle.to_tensor([2, 3], dtype="float32", stop_gradient=True)
-        y = paddle.jit.to_static(net)(x)
-        assert os.path.exists("/tmp/SIR_0.py")
+        with with_export_guard(temp_dir_name):
+            y = paddle.jit.to_static(net)(x)
+        assert os.path.exists(os.path.join(temp_dir_name, "SIR_0.py"))
+        temp_dir.cleanup()
 
 
 if __name__ == "__main__":
