Optimize the definition of decompose config and fix the config saving of test_target_device.

Author: Xreki

graph_net/subgraph_decompose_and_evaluation_step.py

@@ -123,17 +123,16 @@ def save_decompose_config(
 ):
     """Saves the current state to a JSON file."""
 
-    active_models_map = {}
-    split_positions_map = {}
+    tasks_map_copy = {}
     for model_name, task_info in tasks_map.items():
-        active_models_map[model_name] = task_info["original_path"]
-        split_positions_map[model_name] = task_info["split_positions"]
+        tasks_map_copy[model_name] = {}
+        for key in ["original_path", "split_positions"]:
+            tasks_map_copy[model_name][key] = task_info[key]
 
     config = {
         "max_subgraph_size": max_subgraph_size,
         "incorrect_models": list(incorrect_paths),
-        "active_models_map": active_models_map,
-        "split_positions_map": split_positions_map,
+        "tasks_map": tasks_map_copy,
         "failed_decomposition_models": list(failed_decomposition_models),
     }
     config_path = get_decompose_config_path(work_dir)
@@ -283,9 +282,8 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
     print(f"[Init] Resuming from Pass_{current_pass_id - 1} (Dir: {prev_pass_dir})...")
 
     prev_config = load_decompose_config(prev_pass_dir)
-    prev_active_models_map = prev_config.get("active_models_map", {})
-    prev_split_positions_map = prev_config.get("split_positions_map", {})
     prev_incorrect_subgraphs = prev_config.get("incorrect_models", [])
+    prev_tasks_map = prev_config.get("tasks_map", {})
 
     # Load previous max size as fallback
     prev_max_subgraph_size = prev_config.get("max_subgraph_size")
@@ -302,12 +300,12 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
         model_name_with_subgraph_idx = subgraph_path.rstrip("/").split(os.sep)[-1]
         model_name = "_".join(model_name_with_subgraph_idx.split("_")[:-1])
         subgraph_idx = int(model_name_with_subgraph_idx.split("_")[-1])
-        print(f"subgraph_path: {subgraph_path}")
-        print(f"model_name: {model_name}, subgraph_idx: {subgraph_idx}")
-        assert model_name in prev_active_models_map
+
+        assert model_name in prev_tasks_map
+        pre_task_for_model = prev_tasks_map[model_name]
 
         # Reconstruct previous subgraph size to locate the failing segment
-        prev_split_positions = prev_split_positions_map.get(model_name, [])
+        prev_split_positions = pre_task_for_model.get("split_positions", [])
         subgraph_size = reconstruct_subgraph_size(prev_split_positions)
         assert subgraph_idx < len(
             subgraph_size
@@ -316,7 +314,7 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
         if model_name not in tasks_map:
             tasks_map[model_name] = {
                 "subgraph_path": subgraph_path,
-                "original_path": prev_active_models_map[model_name],
+                "original_path": pre_task_for_model["original_path"],
                 "subgraph_size": subgraph_size[subgraph_idx],
                 "split_positions": set(),
             }
@@ -447,6 +445,11 @@ def main(args):
         ) = execute_decomposition_phase(
             max_subgraph_size, tasks_map, args.framework, pass_work_dir
         )
+    else:
+        config = load_decompose_config(pass_work_dir)
+        max_subgraph_size = config["max_subgraph_size"]
+        failed_decomposition = config["failed_decomposition_models"]
+        tasks_map = config.get("tasks_map", {})
 
     # --- Step 4: Testing ---
     pass_log_path = os.path.join(pass_work_dir, "batch_test_result.log")