Merge branch 'opt_saved_results' into add_original_names

Author: Xreki

graph_net/subgraph_decompose_and_evaluation_step.py

@@ -244,6 +244,12 @@ def run_decomposer_for_multi_models(
         original_path = task_info["original_path"]
 
         split_positions = sorted(list(task_info["split_positions"]))
+
+        method = "fixed-start"
+        if method == "fixed-start":
+            assert len(split_positions) >= 3, f"{split_positions=}"
+            split_positions = [0, split_positions[1]]
+
         rectified_model_path = get_rectfied_model_path(original_path)
         assert os.path.exists(
             rectified_model_path
@@ -293,27 +299,23 @@ def run_evaluation(
     ), f"[ERROR] test failed for {work_dir}, please check the log."
 
 
-def reconstruct_subgraph_size(split_positions: List[int]) -> List[list]:
-    """Reconstructs subgraph size based on sorted split positions."""
-    deduplicated_splits = sorted(set(split_positions))
-
-    subgraph_size = [
-        deduplicated_splits[i : i + 2] for i in range(len(deduplicated_splits) - 1)
-    ]
-    return subgraph_size
+def reconstruct_split_positions_for_subgraphs(
+    split_positions, subgraph_idxs, max_subgraph_size
+):
+    subgraph_idxs = [subgraph_idxs] if isinstance(subgraph_idxs, int) else subgraph_idxs
 
+    new_split_positions = []
+    for subgraph_idx in subgraph_idxs:
+        assert (
+            subgraph_idx < len(split_positions) - 1
+        ), f"subgraph_idx {subgraph_idx} is out of bounds of split_positions: {split_positions}."
 
-def calculate_split_positions_for_subgraph(subgraph_range, max_subgraph_size):
-    assert isinstance(subgraph_range, (list, tuple)) and len(subgraph_range) == 2
+        start_pos, end_pos = split_positions[subgraph_idx : subgraph_idx + 2]
+        new_split_positions = new_split_positions + list(
+            range(start_pos, end_pos + max_subgraph_size - 1, max_subgraph_size)
+        )
 
-    # subgraph_size: the start and end position in original model.
-    start_pos, end_pos = subgraph_range
-    end_pos = kMaxGraphSize if end_pos == float("inf") else end_pos
-
-    split_positions = set(
-        range(start_pos, end_pos + max_subgraph_size - 1, max_subgraph_size)
-    )
-    return list(sorted(split_positions))
+    return sorted(list(set(new_split_positions)))
 
 
 def generate_initial_tasks(args):
@@ -322,19 +324,16 @@ def generate_initial_tasks(args):
     initial_failures = get_ranged_incorrect_models(args.tolerance, args.log_file)
 
     tasks_map = {}
-    max_subgraph_size = args.max_subgraph_size
+    max_subgraph_size = min(args.max_subgraph_size, kMaxGraphSize // 2)
 
+    initial_split_positions = reconstruct_split_positions_for_subgraphs(
+        [0, kMaxGraphSize], 0, max_subgraph_size
+    )
     for model_path in initial_failures:
         model_name = get_model_name_with_subgraph_tag(model_path)
-
-        initial_range = [0, kMaxGraphSize]
-        initial_splits = calculate_split_positions_for_subgraph(
-            initial_range, max_subgraph_size
-        )
-
         tasks_map[model_name] = {
             "original_path": model_path,
-            "split_positions": initial_splits,
+            "split_positions": initial_split_positions,
         }
 
     running_states = {
@@ -354,7 +353,29 @@ def extract_model_name_and_subgraph_idx(subgraph_path):
     return model_name, subgraph_idx
 
 
-def generate_successor_tasks(base_output_dir, current_pass_id):
+def collect_incorrect_subgraph_idxs(args, model_names, incorrect_models):
+    model_name2subgraph_idxs = {}
+    for subgraph_path in sorted(incorrect_models):
+        model_name, subgraph_idx = extract_model_name_and_subgraph_idx(subgraph_path)
+        print(f"{subgraph_path=}")
+
+        if model_name not in model_name2subgraph_idxs:
+            model_name2subgraph_idxs[model_name] = []
+        model_name2subgraph_idxs[model_name].append(subgraph_idx)
+
+    if args.method == "fixed-start":
+        for model_name in model_names:
+            if model_name not in model_name2subgraph_idxs:
+                model_name2subgraph_idxs[model_name] = [1]
+            else:
+                assert (
+                    len(model_name2subgraph_idxs[model_name]) == 1
+                    and model_name2subgraph_idxs[model_name] == 0
+                )
+    return model_name2subgraph_idxs
+
+
+def generate_successor_tasks(args, 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)
     print(f"[Init] Resuming from Pass_{current_pass_id - 1} (Dir: {prev_pass_dir})...")
@@ -367,34 +388,24 @@ def generate_successor_tasks(base_output_dir, current_pass_id):
     tasks_map = {}
     prev_tasks_map = prev_config.tasks_map
 
-    for subgraph_path in sorted(prev_config.incorrect_models):
-        model_name, subgraph_idx = extract_model_name_and_subgraph_idx(subgraph_path)
+    model_name2subgraph_idxs = collect_incorrect_subgraph_idxs(
+        args, list(prev_tasks_map.keys()), prev_config.incorrect_models
+    )
 
+    for model_name, subgraph_idxs in model_name2subgraph_idxs.items():
         assert model_name in prev_tasks_map
         pre_task_for_model = prev_tasks_map[model_name]
 
         prev_split_positions = pre_task_for_model.get("split_positions", [])
-        subgraph_ranges = reconstruct_subgraph_size(prev_split_positions)
-
-        assert subgraph_idx < len(
-            subgraph_ranges
-        ), f"subgraph_idx {subgraph_idx} is out of bounds for {model_name} (previous split_positions: {prev_split_positions})"
-
-        split_positions = calculate_split_positions_for_subgraph(
-            subgraph_ranges[subgraph_idx], max_subgraph_size
+        split_positions = reconstruct_split_positions_for_subgraphs(
+            prev_split_positions, subgraph_idxs, max_subgraph_size
         )
-        if model_name not in tasks_map:
-            tasks_map[model_name] = {
-                "original_path": pre_task_for_model["original_path"],
-                "split_positions": list(sorted(split_positions)),
-            }
-        else:
-            merged_split_positions = (
-                tasks_map[model_name]["split_positions"] + split_positions
-            )
-            tasks_map[model_name]["split_positions"] = list(
-                sorted(set(merged_split_positions))
-            )
+
+        tasks_map[model_name] = {
+            "original_path": pre_task_for_model["original_path"],
+            "split_positions": split_positions,
+        }
+        print(f"{tasks_map=}")
 
     return tasks_map, max_subgraph_size, prev_config.running_states
 
@@ -404,7 +415,7 @@ def prepare_tasks_and_verify(args, current_pass_id, base_output_dir):
         tasks_map, max_subgraph_size, running_states = generate_initial_tasks(args)
     else:
         tasks_map, max_subgraph_size, running_states = generate_successor_tasks(
-            base_output_dir, current_pass_id
+            args, base_output_dir, current_pass_id
         )
 
     print(f"[Init] initial max_subgraph_size: {max_subgraph_size}")
@@ -431,6 +442,7 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, workspa
 
     failed_decomposition = []
     need_decompose = True if len(tasks_map) > 0 else False
+    method = "fixed-start"
 
     while need_decompose:
         decomposed_samples_dir = os.path.join(
@@ -455,6 +467,7 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, workspa
             not failed_decomposition
             and num_decomposed_samples == len(tasks_map)
             and max_subgraph_size > 1
+            and method != "fixed-start"
         ):
             need_decompose = True
             shutil.rmtree(decomposed_samples_dir)
@@ -464,8 +477,8 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, workspa
                 split_positions = task_info["split_positions"]
                 if not split_positions or len(split_positions) < 2:
                     continue
-                new_split_positions = calculate_split_positions_for_subgraph(
-                    split_positions[0:2], max_subgraph_size
+                new_split_positions = reconstruct_split_positions_for_subgraphs(
+                    split_positions, 0, max_subgraph_size
                 )
                 task_info["split_positions"] = new_split_positions
         else:
@@ -579,6 +592,7 @@ def main(args):
     parser.add_argument(
         "--test-config", type=str, required=True, help="Base64 encoded test config"
     )
+    parser.add_argument("--method", type=str, required=True)
     parser.add_argument(
         "--tolerance",
         type=int,