Temporally support to save the random states.

Author: Xreki

graph_net/paddle/random_util.py

@@ -0,0 +1,40 @@
+import os
+import pickle
+import numpy as np
+import random
+import re
+import paddle
+
+
+def set_seed(random_seed):
+    paddle.seed(random_seed)
+    random.seed(random_seed)
+    np.random.seed(random_seed)
+
+
+def _generate_random_state_filename(model_path):
+    fields = model_path.rstrip("/").split(os.sep)
+    pattern = r"^subgraph(_\d+)?$"
+    return f"{fields[-2]}_{fields[-1]}" if re.match(pattern, fields[-1]) else fields[-1]
+
+
+def save_random_states(model_path, output_dir, random_state_dict):
+    filepath = os.path.join(output_dir, _generate_random_state_filename(model_path))
+    print(f"Write to {filepath}.", flush=True)
+    try:
+        with open(filepath, "wb") as f:
+            pickle.dump(random_state_dict, f)
+    except Exception:
+        print(f"Fail to open {filepath}.")
+
+
+def load_random_states(model_path, output_dir):
+    filepath = os.path.join(output_dir, _generate_random_state_filename(model_path))
+    print(f"Read from {filepath}.", flush=True)
+    random_states = None
+    try:
+        with open(filepath, "rb") as f:
+            random_states = pickle.load(f)
+    except Exception:
+        print(f"Fail to open {filepath}.")
+    return random_states

graph_net/paddle/run_model.py

@@ -3,19 +3,12 @@
 import base64
 import argparse
 import numpy as np
-import random
 
 os.environ["FLAGS_logging_pir_py_code_dir"] = "/tmp/dump"
 
 import paddle
 from graph_net import imp_util
-from graph_net.paddle import utils
-
-
-def set_seed(random_seed):
-    paddle.seed(random_seed)
-    random.seed(random_seed)
-    np.random.seed(random_seed)
+from graph_net.paddle import utils, random_util
 
 
 def load_class_from_file(file_path: str, class_name: str):
@@ -31,13 +24,20 @@ def get_input_dict(model_path):
     params = inputs_params["weight_info"]
     inputs = inputs_params["input_info"]
 
-    state_dict = {}
-    for k, v in params.items():
-        name = v["original_name"] if v.get("original_name", None) else k
-        state_dict[k] = paddle.nn.parameter.Parameter(utils.replay_tensor(v), name=name)
-    for k, v in inputs.items():
-        state_dict[k] = utils.replay_tensor(v)
-    return state_dict
+    random_state_dict = {}
+    input_dict = {}
+    for name, meta in params.items():
+        original_name = (
+            meta["original_name"] if meta.get("original_name", None) else name
+        )
+        random_state_dict[name] = np.random.get_state()
+        input_dict[name] = paddle.nn.parameter.Parameter(
+            utils.replay_tensor(meta), name=original_name
+        )
+    for name, meta in inputs.items():
+        random_state_dict[name] = np.random.get_state()
+        input_dict[name] = utils.replay_tensor(meta)
+    return input_dict, random_state_dict
 
 
 def _convert_to_dict(config_str):
@@ -71,9 +71,11 @@ def main(args):
     print(f"{model_path=}")
 
     initalize_seed = 123
-    set_seed(random_seed=initalize_seed)
+    random_util.set_seed(random_seed=initalize_seed)
 
-    input_dict = get_input_dict(args.model_path)
+    input_dict, random_state_dict = get_input_dict(args.model_path)
+    output_dir = "/work/GraphNet/graph_net/test/outputs/pass_0"
+    random_util.save_random_states(model_path, output_dir, random_state_dict)
     model = _get_decorator(args)(model)
     model(**input_dict)
 

graph_net/paddle/test_compiler.py

@@ -4,21 +4,15 @@
 from pathlib import Path
 import sys
 import os
-from dataclasses import dataclass
-from contextlib import contextmanager
-import time
-import math
 import numpy as np
-import random
 import platform
 import traceback
 import subprocess
 import re
 
-from graph_net.paddle import utils
 from graph_net import path_utils
 from graph_net import test_compiler_util
-
+from graph_net.paddle import utils, random_util
 from graph_net.paddle.backend.graph_compiler_backend import GraphCompilerBackend
 from graph_net.paddle.backend.cinn_backend import CinnBackend
 from graph_net.paddle.backend.nope_backend import NopeBackend
@@ -35,12 +29,6 @@ def get_compiler_backend(args) -> GraphCompilerBackend:
     return registry_backend[args.compiler]
 
 
-def set_seed(random_seed):
-    paddle.seed(random_seed)
-    random.seed(random_seed)
-    np.random.seed(random_seed)
-
-
 def init_env(args):
     if test_compiler_util.is_gpu_device(args.device):
         paddle.set_flags({"FLAGS_cudnn_exhaustive_search": 1})
@@ -62,7 +50,7 @@ def get_hardward_name(args):
                     )
                 )
             )
-        except Exception as e:
+        except Exception:
             pass
     elif args.device == "cpu":
         hardware = platform.processor()
@@ -104,10 +92,18 @@ def get_input_dict(model_path):
     inputs_params = utils.load_converted_from_text(f"{model_path}")
     params = inputs_params["weight_info"]
     inputs = inputs_params["input_info"]
-
     params.update(inputs)
-    state_dict = {k: utils.replay_tensor(v) for k, v in params.items()}
-    return state_dict
+
+    output_dir = "/work/GraphNet/graph_net/test/outputs/pass_0"
+    random_states = random_util.load_random_states(model_path, output_dir)
+
+    input_dict = {}
+    for name, meta in params.items():
+        if random_states is not None and random_states.get(name, None) is not None:
+            np.random.set_state(random_states[name])
+        tensor = utils.replay_tensor(meta)
+        input_dict[name] = tensor
+    return input_dict
 
 
 def get_input_spec(model_path):
@@ -476,7 +472,7 @@ def main(args):
     assert args.device in ["cuda", "dcu", "xpu", "cpu"]
 
     initalize_seed = 123
-    set_seed(random_seed=initalize_seed)
+    random_util.set_seed(random_seed=initalize_seed)
 
     if path_utils.is_single_model_dir(args.model_path):
         test_single_model(args)

graph_net/paddle/test_reference_device.py

@@ -1,20 +1,15 @@
 import argparse
-import importlib.util
 import paddle
-import time
-import numpy as np
-import random
 import os
 from pathlib import Path
 from contextlib import redirect_stdout, redirect_stderr
 import json
-import re
 import sys
 import traceback
 
 from graph_net import path_utils
 from graph_net import test_compiler_util
-from graph_net.paddle import test_compiler
+from graph_net.paddle import random_util, test_compiler
 
 
 def get_reference_log_path(reference_dir, model_path):
@@ -130,7 +125,7 @@ def main(args):
     assert args.compiler in {"cinn", "nope"}
     assert args.device in ["cuda"]
 
-    test_compiler.set_seed(random_seed=args.seed)
+    random_util.set_seed(random_seed=args.seed)
     test_compiler.init_env(args)
 
     ref_dump_dir = Path(args.reference_dir)