Temporally support to save the random states.

Author: Xreki

graph_net/paddle/run_model.py

@@ -4,6 +4,7 @@
 import argparse
 import numpy as np
 import random
+import pickle
 
 os.environ["FLAGS_logging_pir_py_code_dir"] = "/tmp/dump"
 
@@ -31,13 +32,27 @@ 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 save_random_states(output_dir, random_state_dict):
+    filepath = os.path.join(output_dir, "random_states.pkl")
+    print(f"Write to {filepath}.")
+    with open(filepath, "wb") as f:
+        pickle.dump(random_state_dict, f)
 
 
 def _convert_to_dict(config_str):
@@ -73,7 +88,9 @@ def main(args):
     initalize_seed = 123
     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"
+    save_random_states(output_dir, random_state_dict)
     model = _get_decorator(args)(model)
     model(**input_dict)
 

graph_net/paddle/test_compiler.py

@@ -4,12 +4,9 @@
 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 pickle
 import platform
 import traceback
 import subprocess
@@ -62,7 +59,7 @@ def get_hardward_name(args):
                     )
                 )
             )
-        except Exception as e:
+        except Exception:
             pass
     elif args.device == "cpu":
         hardware = platform.processor()
@@ -100,14 +97,31 @@ def get_model(model_path):
     return model_class()
 
 
+def load_random_states(output_dir):
+    filepath = os.path.join(output_dir, "random_states.pkl")
+    print(f"Read from {filepath}.")
+    random_states = None
+    with open(filepath, "rb") as f:
+        random_states = pickle.load(f)
+    return random_states
+
+
 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 = load_random_states(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):