Record the number of original incorrect models.

Author: Xreki

graph_net/subgraph_decompose_and_evaluation_step.py

@@ -326,6 +326,14 @@ def generate_initial_tasks(args):
     return tasks_map, max_subgraph_size, running_states
 
 
+def extract_model_name_and_subgraph_idx(subgraph_path):
+    # 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])
+    return model_name, subgraph_idx
+
+
 def generate_refined_tasks(base_output_dir, current_pass_id):
     """Generates tasks for Pass > 0 based on previous pass results."""
     prev_pass_dir = get_decompose_workspace_path(base_output_dir, current_pass_id - 1)
@@ -340,10 +348,7 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
     prev_tasks_map = prev_config.tasks_map
 
     for subgraph_path in sorted(prev_config.incorrect_models):
-        # 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])
+        model_name, subgraph_idx = extract_model_name_and_subgraph_idx(subgraph_path)
 
         assert model_name in prev_tasks_map
         pre_task_for_model = prev_tasks_map[model_name]
@@ -525,12 +530,18 @@ def main(args):
         )
         print(f"\n--- Phase 3: Analysis (torlance={tolerance}) ---")
         next_round_models = sorted(get_incorrect_models(tolerance, pass_log_path))
+        original_model_paths = [
+            model_name
+            for subgraph_path in next_round_models
+            for model_name, _ in [extract_model_name_and_subgraph_idx(subgraph_path)]
+        ]
+
         running_states[f"pass_{current_pass_id + 1}"] = {
-            "num_incorrect_models": len(next_round_models),
+            "num_incorrect_models": len(set(original_model_paths)),
             "incorrect_models": list(next_round_models),
         }
 
-        print(f"[Analysis] Found {len(next_round_models)} incorrect subgraphs.\n")
+        print(f"[Analysis] Found {len(next_round_models)} incorrect subgraphs.")
         for idx, model_path in enumerate(next_round_models):
             print(f"- [{idx}] {model_path}")