Fix the replay of tensor, not use dummy in decomposer. (#454)

Author: Xreki

graph_net/torch/fx_graph_module_util.py

@@ -5,10 +5,10 @@
 from dataclasses import asdict
 
 
-def get_torch_module_and_inputs(model_path):
+def get_torch_module_and_inputs(model_path, use_dummy_inputs=True):
     module = _get_torch_module(model_path)
     tensor_metas = _get_tensor_metas(model_path)
-    inputs = _create_inputs_by_metas(module, tensor_metas)
+    inputs = _create_inputs_by_metas(module, tensor_metas, use_dummy_inputs)
     return module, inputs
 
 
@@ -27,11 +27,11 @@ def _get_tensor_metas(model_path):
     ]
 
 
-def _create_inputs_by_metas(module, tensor_metas):
+def _create_inputs_by_metas(module, tensor_metas, use_dummy_inputs):
     tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-    from graph_net.torch.utils import get_dummy_named_tensors
+    from graph_net.torch.utils import get_named_tensors
 
-    named_tensors = get_dummy_named_tensors(tensor_meta_attrs_list)
+    named_tensors = get_named_tensors(tensor_meta_attrs_list, use_dummy_inputs)
     name2tensor = {k: v for k, v in named_tensors}
     return tuple(
         name2tensor[name] for name in inspect.signature(module.forward).parameters

graph_net/torch/graph_decomposer.py

@@ -130,7 +130,7 @@ def __call__(self, rel_model_path):
             for k, v in self.config.items()
             if k in {"split_positions", "group_head_and_tail", "chain_style"}
         }
-        module, inputs = get_torch_module_and_inputs(model_path)
+        module, inputs = get_torch_module_and_inputs(model_path, use_dummy_inputs=False)
         gm = parse_immutable_model_path_into_sole_graph_module(model_path)
         try:
             # logger.warning("convert_to_submodules_graph-call-begin")
@@ -227,7 +227,7 @@ def __call__(self, rel_model_path):
             "group_head_and_tail": self.config.get("group_head_and_tail", False),
             "chain_style": self.config.get("chain_style", False),
         }
-        module, inputs = get_torch_module_and_inputs(model_path)
+        module, inputs = get_torch_module_and_inputs(model_path, use_dummy_inputs=False)
         gm = parse_sole_graph_module(module, inputs)
         rewrited_gm: torch.fx.GraphModule = convert_to_submodules_graph(
             gm,

graph_net/torch/utils.py

@@ -236,15 +236,18 @@ def convert_tensor_meta_attrs_list_to_named_tensors(tensor_meta_attrs_list):
     return ret
 
 
-def get_dummy_named_tensors(tensor_meta_attrs_list):
+def get_named_tensors(tensor_meta_attrs_list, use_dummy_inputs):
     tensors_wrappers = convert_tensor_meta_attrs_list_to_tensors_wrappers(
         tensor_meta_attrs_list
     )
     ret = []
     for i, tensors_wrapper in enumerate(tensors_wrappers):
         name = tensors_wrapper["name"]
         # shape = tensors_wrapper["info"]['shape']
-        tensor = get_dummy_tensor(tensors_wrapper)
+        if use_dummy_inputs:
+            tensor = get_dummy_tensor(tensors_wrapper)
+        else:
+            tensor = replay_tensor(tensors_wrapper)
         ret.append((name, tensor))
     return ret