1) fix bugs in decomposer_util.py; 2) support tools/typical_sequence_decompose.sh (#445)

lixinqi · web-flow · commit d6bcceada93e · 2025-12-12T17:44:26.000+08:00
* init 'symbolic_dimension_reifier' field in graph_net.json

* remove unused files

* add DeviceRewriteSamplePass

* fix output node order bug in torch/decompose_util.py

* update graph_hash.txt in OnlyModelFileRewriteSamplePassMixin

* 1) fix bugs in decomposer_util.py; 2) support tools/typical_sequence_decompose.sh
diff --git a/graph_net/tensor_meta.py b/graph_net/tensor_meta.py
@@ -1,6 +1,7 @@
 import importlib.util as imp
 import inspect
 from dataclasses import dataclass
+import math
 
 
 @dataclass
@@ -83,6 +84,8 @@ def serialize_to_py_str(self) -> str:
     def _get_limited_precision_float_str(self, value):
         if not isinstance(value, float):
             return value
+        if math.isnan(value) or math.isinf(value):
+            return f'float("{value}")'
         return f"{value:.3f}"
 
     def _format_data(self, data):
diff --git a/graph_net/tools/typical_sequence_decompose.sh b/graph_net/tools/typical_sequence_decompose.sh
@@ -0,0 +1,81 @@
+#!/bin/bash
+set -x
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
+DECOMPOSE_WORKSPACE=/tmp/typical_sequence_decompose_workspace
+
+mkdir -p "$DECOMPOSE_WORKSPACE"
+
+model_list="$GRAPH_NET_ROOT/graph_net/test/dev_model_list/validation_error_model_list.txt"
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list $model_list \
+    --handler-config=$(base64 -w 0 <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/typical_sequence_split_points.py",
+    "handler_class_name": "OpNamesExtractor",
+    "handler_config": {
+        "resume": true,
+        "model_path_prefix": "$GRAPH_NET_ROOT",
+        "output_dir": "$DECOMPOSE_WORKSPACE"
+    }
+}
+EOF
+)
+
+python3 -m graph_net.torch.typical_sequence_split_points \
+    --enable-resume \
+    --model-list "$model_list" \
+    --op-names-path-prefix "$DECOMPOSE_WORKSPACE" \
+    --device "cuda" \
+    --window-size 10 \
+    --fold-policy default \
+    --fold-times 10 \
+    --output-json "$DECOMPOSE_WORKSPACE/split_results.json"
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list $model_list \
+    --handler-config=$(base64 -w 0 <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/graph_decomposer.py",
+    "handler_class_name": "RangeDecomposerExtractor",
+    "handler_config": {
+        "resume": true,
+        "model_path_prefix": "$GRAPH_NET_ROOT",
+        "output_dir": "$DECOMPOSE_WORKSPACE",
+        "split_results_path": "$DECOMPOSE_WORKSPACE/split_results.json",
+        "group_head_and_tail": true,
+        "chain_style": false
+    }
+}
+EOF
+)
+
+subgraph_sample_list=$DECOMPOSE_WORKSPACE/subgraph_sample_list.txt
+cat $model_list \
+    | grep -v '# ' \
+    | xargs -I {} find $DECOMPOSE_WORKSPACE/{} -name "model.py" \
+    | xargs dirname \
+    | xargs realpath --relative-to=$DECOMPOSE_WORKSPACE \
+    | tee $subgraph_sample_list
+
+GRAPH_VAR_RENAME_WORKSPACE=$DECOMPOSE_WORKSPACE/graph_var_renamed
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list $subgraph_sample_list \
+    --handler-config=$(base64 -w 0 <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/graph_variable_renamer.py",
+    "handler_class_name": "GraphVariableRenamer",
+    "handler_config": {
+        "model_path_prefix": "$DECOMPOSE_WORKSPACE",
+        "data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
+        "data_input_predicator_class_name": "NaiveDataInputPredicator",
+        "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
+        "model_runnable_predicator_class_name": "ModelRunnablePredicator",
+        "output_dir": "$GRAPH_VAR_RENAME_WORKSPACE"
+    }
+}
+EOF
+)
+
diff --git a/graph_net/torch/decompose_util.py b/graph_net/torch/decompose_util.py
@@ -85,7 +85,8 @@ def print_submodule_call(prompt, gm):
     def sort_key(node):
         return new_node2original_node[node].name
 
-    for range_idx in range(len(range_idx2submodule_body_nodes)):
+    num_subgraphs = len(range_idx2submodule_body_nodes)
+    for range_idx in range(num_subgraphs):
         (
             submodule_input_nodes,
             submodule_output_nodes,
@@ -96,6 +97,7 @@ def sort_key(node):
             end_node_idx=get_end_node_idx(range_idx),
             chain_style=chain_style,
         )
+        identity_node_set = set(identity_nodes)
 
         def get_input_nodes(range_idx):
             return sorted(submodule_input_nodes, key=sort_key)
@@ -153,13 +155,13 @@ def get_output_nodes(range_idx):
                 prev_node = new_output_node
 
         # Replace all use of outputs
-        identity_node_set = set(identity_nodes)
         for original_output in get_output_nodes(range_idx):
-            if original_output not in identity_node_set:
-                original_output.replace_all_uses_with(node_map[original_output])
-                new_node2original_node[
-                    node_map[original_output]
-                ] = new_node2original_node[original_output]
+            if original_output in identity_node_set:
+                continue
+            original_output.replace_all_uses_with(node_map[original_output])
+            new_node2original_node[node_map[original_output]] = new_node2original_node[
+                original_output
+            ]
 
         # Erase old nodes
         for node in reversed(get_body_nodes(range_idx)):
@@ -215,12 +217,18 @@ def _get_submodule_inputs_and_outputs(
         return minimal_input_nodes, minimal_output_nodes, []
     else:
         node_list = list(gm.graph.nodes)
-        input_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
-            gm=gm, start_node_idx=start_node_idx, end_node_idx=len(node_list)
-        )
-        output_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
-            gm=gm, start_node_idx=end_node_idx, end_node_idx=len(node_list)
-        )
+        if _is_node_idx_out_of_range(gm, start_node_idx):
+            input_nodes = list(_get_return_nodes(gm))
+        else:
+            input_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
+                gm=gm, start_node_idx=start_node_idx, end_node_idx=len(node_list)
+            )
+        if _is_node_idx_out_of_range(gm, end_node_idx):
+            output_nodes = list(_get_return_nodes(gm))
+        else:
+            output_nodes, _ = _get_minimal_submodule_inputs_and_outputs(
+                gm=gm, start_node_idx=end_node_idx, end_node_idx=len(node_list)
+            )
         identity_nodes_set = set(input_nodes) & set(output_nodes)
         identity_nodes = [node for node in input_nodes if node in identity_nodes_set]
         return input_nodes, output_nodes, identity_nodes
@@ -275,25 +283,19 @@ def get_args_node_and_self_node(node):
         for related_node in get_args_node_and_self_node(node):
             count_ctx.node2after_output[related_node] += 1
 
-    if _is_node_idx_out_of_range(gm, start_node_idx):
-        input_nodes = list(_get_return_nodes(gm))
-    else:
-        input_nodes = [
-            node
-            for node in node_list
-            if count_ctx.node2before_input[node] > 0
-            if count_ctx.node2body[node] > 0
-        ]
-    if _is_node_idx_out_of_range(gm, end_node_idx):
-        output_nodes = list(_get_return_nodes(gm))
-    else:
-        output_nodes = [
-            node
-            for node in node_list
-            if not (count_ctx.node2before_input[node] > 0)
-            if count_ctx.node2body[node] > 0
-            if count_ctx.node2after_output[node] > 0
-        ]
+    input_nodes = [
+        node
+        for node in node_list
+        if count_ctx.node2before_input[node] > 0
+        if count_ctx.node2body[node] > 0
+    ]
+    output_nodes = [
+        node
+        for node in node_list
+        if not (count_ctx.node2before_input[node] > 0)
+        if count_ctx.node2body[node] > 0
+        if count_ctx.node2after_output[node] > 0
+    ]
     return input_nodes, output_nodes
 
 
diff --git a/graph_net/torch/graph_variable_renamer.py b/graph_net/torch/graph_variable_renamer.py
@@ -79,10 +79,8 @@ def __call__(self, rel_model_path):
         module, inputs = get_torch_module_and_inputs(src_model_path)
         gm = parse_sole_graph_module(module, inputs)
         gm = self.rename_graph_variables(gm, inputs, src_model_path)
-        model_name = os.path.basename(rel_model_path.rstrip(os.sep))
-        new_rel_path = f"{model_name}_renamed"
         dst_model_path = os.path.realpath(
-            os.path.join(self.config["output_dir"], new_rel_path)
+            os.path.join(self.config["output_dir"], rel_model_path)
         )
         Path(dst_model_path).parent.mkdir(parents=True, exist_ok=True)
         shutil.copytree(src_model_path, dst_model_path, dirs_exist_ok=True)
