merge

Author: JewelRoam

graph_net/test/dev_model_list/validation_error_model_list.txt

@@ -1,27 +1,20 @@
-samples/ultralytics/yolov3-tinyu
-samples/torchgeometric/RECT_L
-samples/transformers-auto-model/opus-mt-en-gv
 samples/timm/convnextv2_base.fcmae_ft_in1k
-samples/torchvision/vgg16_bn
-samples/timm/regnety_080_tv.tv2_in1k
+samples/timm/hgnet_tiny.paddle_in1k
 samples/timm/mobilenetv4_conv_aa_large.e230_r384_in12k
+samples/timm/regnety_080_tv.tv2_in1k
+samples/timm/res2net50_14w_8s.in1k
+samples/torchaudio/wavlm_base
+samples/torchgeometric/RECT_L
+samples/torchvision/vgg16_bn
+samples/transformers-auto-model/bge-small-en-v1.5
 samples/transformers-auto-model/distilbert_distilbert-base-multilingual-cased
 samples/transformers-auto-model/OFA-Sys_chinese-clip-vit-large-patch14
-samples/transformers-auto-model/opus-mt-tc-bible-big-deu_eng_fra_por_spa-bat
-samples/transformers-auto-model/opus-mt-en-tw
-samples/timm/hgnet_tiny.paddle_in1k
 samples/transformers-auto-model/opus-mt-ase-es
-samples/timm/resnetv2_18.ra4e3600r224_in1k
-samples/transformers-auto-model/TinyLlama_TinyLlama-1.1B-Chat-v0.4
-samples/timm/resnetaa50d.din12k
-samples/transformers-auto-model/bge-small-en-v1.5
-samples/timm/res2net50_14w_8s.in1k
-samples/torchaudio/wavlm_base
-samples/transformers-auto-model/opus-mt-en-sal
+samples/transformers-auto-model/opus-mt-en-gv
 samples/transformers-auto-model/opus-mt-en-phi
+samples/transformers-auto-model/opus-mt-en-sal
+samples/transformers-auto-model/opus-mt-en-tw
 samples/transformers-auto-model/opus-mt-fi-niu
+samples/transformers-auto-model/opus-mt-tc-bible-big-deu_eng_fra_por_spa-bat
 samples/transformers-auto-model/opus-mt-tc-bible-big-gmw-deu_eng_fra_por_spa
-samples/transformers-auto-model/opus-mt-NORTHEU-NORTHEU
-samples/transformers-auto-model/sentence-transformers_paraphrase-distilroberta-base-v1
-samples/transformers-auto-model/TrustSafeAI_RADARVicuna7B
-
+samples/ultralytics/yolov3-tinyu

graph_net/test/naive_decomposer_and_post_extract_process_test.sh

@@ -1,5 +1,4 @@
 #!/bin/bash
-# bash graph_net/test/naive_decomposer_and_post_extract_process_test.sh
 
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")

graph_net/test/typical_sequence_decomposer_test.sh

@@ -57,4 +57,4 @@ python3 -m graph_net.torch.test_compiler \
 
 python3 -m graph_net.plot_ESt \
     --benchmark-path "$DECOMPOSE_PATH/validation.log" \
-    --output-dir "$DECOMPOSE_PATH"
+    --output-dir "$DECOMPOSE_PATH"

graph_net/torch/fx_graph_parse_util.py

@@ -122,19 +122,28 @@ def _rename_placeholder(name, pattern2replacement):
     return name
 
 
-def parse_sole_graph_module(module, inputs):
+def parse_sole_graph_module_without_varify(module, inputs):
     traced_module = None
     traced_sample_inputs = None
 
     def my_backend(gm, sample_inputs):
         nonlocal traced_module
-        traced_module = gm
         nonlocal traced_sample_inputs
+        assert traced_module is None
+        assert traced_sample_inputs is None
+        traced_module = gm
         traced_sample_inputs = sample_inputs
         return gm.forward
 
     torch.compile(module, backend=my_backend)(*inputs)
     assert traced_module is not None
+    return traced_module, traced_sample_inputs
+
+
+def parse_sole_graph_module(module, inputs):
+    traced_module, traced_sample_inputs = parse_sole_graph_module_without_varify(
+        module, inputs
+    )
 
     def get_input_names_from_signature():
         return inspect.signature(module.forward).parameters

graph_net/torch/graph_decomposer.py

@@ -262,7 +262,6 @@ def forward(self, *args):
         if not self.extracted:
             if self.need_extract(self.submodule, args):
                 self.builtin_extractor(self.submodule, args)
-                self._post_extract_process()
             self.extracted = True
         return self.submodule(*args)
 

graph_net/torch/typical_sequence_split_points.py

@@ -3,13 +3,11 @@
 import os
 from pathlib import Path
 from typing import Any, Dict, List
-
 import torch
 import torch.nn as nn
-import tempfile
-import graph_net.imp_util
-from graph_net.torch import utils as graph_utils
 from graph_net.torch.rp_expr.rp_expr_parser import RpExprParser
+from graph_net.torch.fx_graph_module_util import get_torch_module_and_inputs
+from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module_without_varify
 
 
 class TypicalSequenceExtractor:
@@ -28,9 +26,12 @@ def _extract_operators_from_graph(
 
             if node.op == "call_module":
                 target_name = type(named_modules[node.target]).__name__
-            else:
+            elif node.op == "call_method":
+                target_name = f"Tensor.{node.target}"
+            elif node.op == "call_function":
                 target_name = getattr(node.target, "__name__", str(node.target))
-
+            else:
+                raise NotImplementedError()
             operator_list.append(
                 {
                     "op_type": node.op,
@@ -48,39 +49,6 @@ def extract_compiler(self, gm: torch.fx.GraphModule, inputs: List[torch.Tensor])
         return gm.forward
 
 
-class TypicalSequenceModelLoader:
-    def load_class_from_file(self, model_path: str, device: str) -> Any:
-        file_path = os.path.join(model_path, "model.py")
-
-        if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Model file not found: {file_path}")
-
-        with open(file_path, "r", encoding="utf-8") as f:
-            model_code = f.read()
-        model_code = graph_utils.modify_code_by_device(model_code, device)
-
-        with tempfile.NamedTemporaryFile(
-            mode="w", suffix=".py", encoding="utf-8"
-        ) as temp_file:
-            temp_file.write(model_code)
-            module = graph_net.imp_util.load_module(temp_file.name)
-        model_class = getattr(module, "GraphModule", None)
-
-        return model_class
-
-    def get_input_dict(self, model_path: str, device: str) -> Dict[str, torch.Tensor]:
-        inputs_params = graph_utils.load_converted_from_text(f"{model_path}")
-        params = inputs_params["weight_info"]
-        for tensor_meta in params.values():
-            if hasattr(tensor_meta, "device"):
-                tensor_meta.device = device
-        input_dict = {
-            k: graph_utils.replay_tensor(v).to(torch.device(device))
-            for k, v in params.items()
-        }
-        return input_dict
-
-
 class SplitAnalyzer:
     def __init__(
         self, window_size: int = 10, fold_policy: str = "default", fold_times: int = 0
@@ -109,20 +77,11 @@ def _resolve_token_to_ops(
     def _extract_ops_via_compile(
         self, model_path: str, device: str = "cpu"
     ) -> List[str]:
-        loader = TypicalSequenceModelLoader()
-        print(f"Loading model from {model_path} on {device}...")
-        try:
-            model_class = loader.load_class_from_file(model_path, device)
-            model = model_class().to(torch.device(device))
-            model.eval()
-            input_dict = loader.get_input_dict(model_path, device)
-        except Exception as e:
-            print(f"Error loading/preparing model {model_path}: {e}")
-            return []
-
+        print(f"extracting ops from {model_path}")
         extractor = TypicalSequenceExtractor()
-        compiled_model = torch.compile(model, backend=extractor.extract_compiler)
-        compiled_model(**input_dict)
+        model, inputs = get_torch_module_and_inputs(model_path)
+        compiled_model, _ = parse_sole_graph_module_without_varify(model, inputs)
+        extractor.extract_compiler(compiled_model, inputs)
         ops_info = extractor.extract_node
 
         return [op["target_name"] for op in ops_info]
@@ -174,7 +133,6 @@ def analyze(
 
         if not inputs_seqs:
             return {}
-
         rp_parser = RpExprParser(
             window_size=self.window_size,
             fold_policy=self.fold_policy,