remove unnecessary code blocks and variables

Author: roll-away

graph_net/torch/naive_graph_decomposer.py

@@ -92,36 +92,28 @@ def forward(self, *args):
         if not self.extracted:
             if self.need_extract(self.submodule, args):
                 self.builtin_extractor(self.submodule, args)
-                self.get_post_extract_process(self.submodule, args)
+                self.get_post_extract_process()
             self.extracted = True
         return self.submodule(*args)
 
     def need_extract(self, gm, sample_inputs):
-        # print("need_extract")
         if self.filter is None:
             return True
-        # if self.fusionablity_filter is not None:
-        #     print("fusionablity of this model is ", self.fusionablity_filter(gm, sample_inputs))
         return self.filter(gm, sample_inputs)
 
-    def get_post_extract_process(self, gm, sample_inputs):
-        # print("modelname: ",self.modelname)
-        # print("parent_graph_extractor.config: ",self.parent_graph_extractor.config['output_dir'])
-        # print("get_post_extract_process")
+    def get_post_extract_process(self):
         model_path = os.path.join(
             self.parent_graph_extractor.config["output_dir"], self.modelname
         )
         return self.post_extract_process(model_path)
 
     def make_filter(self, config):
-        # print("make_filter")
         if config["filter_path"] is None:
             return None
         module = imp_util.load_module(config["filter_path"])
         return module.GraphFilter(config["filter_config"])
 
     def make_post_extract_process(self, config):
-        # print("make post_extract_process")
         if config["filter_path"] is None:
             return None
         module = imp_util.load_module(config["post_extract_process_path"])

graph_net/torch/naive_subgraph_filter.py

@@ -3,6 +3,5 @@ def __init__(self, config):
         self.config = config
 
     def __call__(self, gm, sample_inputs):
-        print("GraphFilter")
         # print(f"GraphFilter\n{gm.code}")
         return True

graph_net/torch/post_extract_process.py

@@ -20,36 +20,8 @@ def __init__(self, config):
         self.config = config
 
     def __call__(self, model_path=None):
-        print("PostExtractProcess")
         if model_path is None:
             return False
-        import json
-        import base64
-        import sys
-        import os
-
-        json_string = json.dumps(self.config)
-        json_bytes = json_string.encode("utf-8")
-        b64_encoded_bytes = base64.b64encode(json_bytes)
-        decorator_config = b64_encoded_bytes.decode("utf-8")
-
-        # args
-        parser = argparse.ArgumentParser(description="load and run model")
-        parser.add_argument(
-            "--model-path",
-            type=str,
-            required=True,
-            help="Path to folder e.g '../../samples/torch/resnet18'",
-        )
-        parser.add_argument(
-            "--decorator-config",
-            type=str,
-            required=False,
-            default=None,
-            help="decorator configuration string",
-        )
-        args = parser.parse_args()
-
         # model
         model_class = load_class_from_file(
             f"{model_path}/model.py", class_name="GraphModule"
@@ -58,46 +30,20 @@ def __call__(self, model_path=None):
         model = model_class()
         print(f"{model_path=}")
 
-        model = _get_decorator(args)(model)
-
         inputs_params = utils.load_converted_from_text(f"{model_path}")
         params = inputs_params["weight_info"]
         state_dict = {k: utils.replay_tensor(v) for k, v in params.items()}
 
         compiled_num_of_kernels = compile_and_count_kernels(model, state_dict)
-        print("compiled: nums_of_kernels = ", compiled_num_of_kernels)
         if compiled_num_of_kernels == 1:
-            print("Graph is fully fusionable")
+            print(model_path, "can be fully integrated")
             return True
         else:
-            print(f"Graph is not fully fusionable ({compiled_num_of_kernels} kernels)")
+            print(model_path, "can not be fully integrated")
             shutil.rmtree(model_path)
             return False
 
 
-def _convert_to_dict(config_str):
-    if config_str is None:
-        return {}
-    config_str = base64.b64decode(config_str).decode("utf-8")
-    config = json.loads(config_str)
-    assert isinstance(config, dict), f"config should be a dict. {config_str=}"
-    return config
-
-
-def _get_decorator(args):
-    if args.decorator_config is None:
-        return lambda model: model
-    decorator_config = _convert_to_dict(args.decorator_config)
-    if "decorator_path" not in decorator_config:
-        return lambda model: model
-    class_name = decorator_config.get("decorator_class_name", "RunModelDecorator")
-    decorator_class = load_class_from_file(
-        decorator_config["decorator_path"],
-        class_name=class_name,
-    )
-    return decorator_class(decorator_config.get("decorator_config", {}))
-
-
 def load_class_from_file(file_path: str, class_name: str) -> Type[torch.nn.Module]:
     spec = importlib.util.spec_from_file_location("unnamed", file_path)
     unnamed = importlib.util.module_from_spec(spec)
@@ -133,11 +79,9 @@ def compile_and_count_kernels(gm, sample_inputs) -> int:
     ) as prof:
         with record_function("model_inference"):
             output = compiled_gm(**sample_inputs)
-    print(prof.key_averages().table())  # print a table of profiler result
     events = prof.key_averages()
     if_compile_work = any(e.key == "TorchDynamo Cache Lookup" for e in events)
     if not if_compile_work:
-        print("Compile failed")
         return -1
     for e in events:
         if e.key == "cuLaunchKernel":