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")
@@ -294,20 +292,18 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
     if not prev_incorrect_subgraphs:
         return {}, max_subgraph_size
 
-    print("[Analysis] Refining splits based on previous incorrect models ...")
-
     tasks_map = {}
     for subgraph_path in prev_incorrect_subgraphs:
         # Parse model name and subgraph index
         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 +312,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(),
             }
@@ -338,7 +334,7 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, pass_wo
         )
         if not os.path.exists(decomposed_samples_dir):
             os.makedirs(decomposed_samples_dir, exist_ok=True)
-            print(f"decomposed_samples_dir: {decomposed_samples_dir}")
+            print(f"- decomposed_samples_dir: {decomposed_samples_dir}")
 
         for model_name, task_info in tasks_map.items():
             print(f"[Decomposition] max_subgraph_size: {max_subgraph_size}")
@@ -387,8 +383,8 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, pass_wo
     return tasks_map, failed_decomposition, max_subgraph_size
 
 
-def print_final_summary(next_round_models, max_subgraph_size):
-    """Prints the final suggestion/result."""
+def print_summary_and_suggestion(next_round_models, max_subgraph_size):
+    """Print suggestion/result."""
     print("\n" + "=" * 80)
     if next_round_models and max_subgraph_size > 1:
         print(f">>> [SUGGESTION] Issues remain (Count: {len(next_round_models)}).")
@@ -447,6 +443,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")
@@ -470,7 +471,7 @@ def main(args):
         failed_decomposition,
     )
 
-    print_final_summary(next_round_models, max_subgraph_size)
+    print_summary_and_suggestion(next_round_models, max_subgraph_size)
 
 
 if __name__ == "__main__":