Naive chainable decompose minor fix (#345)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* minor fix for naive_graph_decomposer

* bug fix for chain_style

Author: lixinqi

graph_net/test/chain_naive_graph_decomposer_test.sh

@@ -7,8 +7,8 @@ MODEL_PATH_IN_SAMPLES=/timm/resnet18
 read -r -d '' json_str <<'EOF'
 {
     "output_dir": "/tmp/naive_decompose_workspace",
-    "split_positions": [2, 4],
-    "group_head_and_tail": false,
+    "split_positions": [8, 16, 32],
+    "group_head_and_tail": true,
     "chain_style": true
 }
 EOF

graph_net/torch/decompose_util.py

@@ -82,6 +82,7 @@ def print_submodule_call(prompt, gm):
         (
             submodule_input_nodes,
             submodule_output_nodes,
+            identity_nodes,
         ) = _get_submodule_inputs_and_outputs(
             original_gm=original_gm,
             start_node_idx=get_start_node_idx(range_idx),
@@ -141,8 +142,10 @@ 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):
-            original_output.replace_all_uses_with(node_map[original_output])
+            if original_output not in identity_node_set:
+                original_output.replace_all_uses_with(node_map[original_output])
 
         # Erase old nodes
         for node in reversed(get_body_nodes(range_idx)):
@@ -215,33 +218,38 @@ def get_related_node(node):
         for related_node in get_related_node(node):
             count_ctx.node2after_output[related_node] += 1
 
-    if chain_style:
-        input_nodes = [
-            node
-            for node in node_list
-            if (count_ctx.node2before_input[node] > 0)
-            if (count_ctx.node2body[node] > 0 or count_ctx.node2after_output[node] > 0)
-        ]
-        input_nodes_set = set(input_nodes)
-        output_nodes = [
-            node
-            for node in node_list
-            if (count_ctx.node2before_input[node] > 0 or 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
+    ]
+    if not chain_style:
+        identity_nodes = []
     else:
-        input_nodes = [
+        identity_nodes = [
             node
             for node in node_list
             if count_ctx.node2before_input[node] > 0
-            if count_ctx.node2body[node] > 0
+            if count_ctx.node2body[node] == 0
+            if count_ctx.node2after_output[node] > 0
+        ][:1]
+        input_nodes_set = set(input_nodes)
+        input_nodes = [
+            *input_nodes,
+            *[node for node in identity_nodes if node not in input_nodes_set],
         ]
+        output_nodes_set = set(output_nodes)
         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
+            *output_nodes,
+            *[node for node in identity_nodes if node not in output_nodes_set],
         ]
 
-    return input_nodes, output_nodes
+    return input_nodes, output_nodes, identity_nodes