Fix

Author: JewelRoam

graph_net/torch/naive_graph_decomposer.py

@@ -31,8 +31,11 @@ def make_config(
         filter_config=None,
         post_extract_process_path=None,
         post_extract_process_class_name=None,
+        post_extract_process_config=None,
         **kwargs,
     ):
+        if post_extract_process_config is None:
+            post_extract_process_config = {}
         for pos in split_positions:
             assert isinstance(
                 pos, int
@@ -46,6 +49,7 @@ def make_config(
             "filter_config": filter_config if filter_config is not None else {},
             "post_extract_process_path": post_extract_process_path,
             "post_extract_process_class_name": post_extract_process_class_name,
+            "post_extract_process_config": post_extract_process_config,
         }
 
     def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
@@ -112,8 +116,8 @@ def make_filter(self, config):
         return module.GraphFilter(config["filter_config"])
 
     def make_post_extract_process(self, config):
-        if config["post_extract_process_path"] is None:
-            return None
+        if config.get("post_extract_process_path") is None:
+            return lambda *args, **kwargs: None
         module = imp_util.load_module(config["post_extract_process_path"])
         cls = getattr(module, config["post_extract_process_class_name"])
-        return cls(config["post_extract_process_path"])
+        return cls(config["post_extract_process_config"])

graph_net/torch/typical_sequence_split_points.py

@@ -2,7 +2,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List
+from typing import Any, Dict, List
 
 import torch
 import torch.nn as nn
@@ -252,7 +252,15 @@ def _print_analysis(self, name, path, splits, total_len, full_ops):
         print("\n")
 
 
-def main():
+def main(args):
+    analyzer = SplitAnalyzer(window_size=args.window_size)
+    results = analyzer.analyze(args.model_list, args.device)
+    if args.output_json:
+        with open(args.output_json, "w") as f:
+            json.dump(results, f, indent=4)
+
+
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Analyze graph and calculate split points."
     )
@@ -278,14 +286,4 @@ def main():
         help="Path to save the analysis results in JSON format.",
     )
     args = parser.parse_args()
-
-    analyzer = SplitAnalyzer(window_size=args.window_size)
-    results = analyzer.analyze(args.model_list, args.device)
-
-    if args.output_json:
-        with open(args.output_json, "w") as f:
-            json.dump(results, f, indent=4)
-
-
-if __name__ == "__main__":
-    main()
+    main(args)