Parse logs

benchmarker/__main__.py

@@ -37,6 +37,23 @@ def run_cmd_and_get_output(command):
     result = filter_json_from_output(process_out)
     return result
 
+def parse_dict(profile_dict):
+    for key,value in profile_dict.items():
+        if isinstance(value, dict):
+            parse_dict(value)
+        elif key == "param":
+            params = {}
+            if "Linear" in value:
+                params["name"] = "fc"
+            else:
+                params["name"] = value.split('(')[0].lower()
+            params["args"] = value
+            profile_dict[key] = params
+
+    for key,value in profile_dict.items():
+        if key == "null" and isinstance(value["param"],dict):
+            profile_dict[value["param"]["name"]] = profile_dict.pop(key)
+
 def main():
     parser = argparse.ArgumentParser(description="Benchmark me up, Scotty!")
     parser.add_argument("--flops", action="store_true")
@@ -73,6 +90,7 @@ def main():
         result["profile_pytorch"] = True
         result["profile_data"] = profile_result["profile_data"]
         result["path_out"] = "./logs/profile"
+        parse_dict(result["profile_data"])
 
     cute_device = get_cute_device_str(result["device"]).replace(" ", "_")
     result["path_out"] = os.path.join(result["path_out"], result["problem"]["name"])

benchmarker/modules/torchprof.py

@@ -6,8 +6,7 @@
 from collections import namedtuple, defaultdict, OrderedDict
 
 Trace = namedtuple("Trace", ["path", "leaf", "module"])
-Measure = namedtuple("Measure", ["self_cpu_total", "cpu_total", "cuda_total", "occurrences", "param"])
-#Measure = namedtuple("Measure", ["self_cpu_total", "cpu_total", "cuda_total", "occurrences"])
+Measure = namedtuple("Measure", ["self_cpu_total", "cpu_total", "cuda_total", "occurrences", "param", "input_shape"])
 
 def walk_modules(module, name="", path=()):
     """Generator. Walks through a PyTorch Module and outputs Trace tuples"""
@@ -73,7 +72,7 @@ def _hook_trace(self, trace):
 
             @functools.wraps(_forward)
             def wrap_forward(*args, **kwargs):
-                with tprofiler.profile(use_cuda=self.use_cuda) as prof:
+                with tprofiler.profile(use_cuda=self.use_cuda,record_shapes=True) as prof:
                     res = _forward(*args, **kwargs)
                 event_list = prof.function_events
                 event_list.populate_cpu_children()
@@ -124,7 +123,6 @@ def traces_to_display(traces, trace_events, show_events=False, paths=None):
         current_tree = tree
         # unwrap all of the events, in case model is called multiple times
         events = [te for tevents in trace_events[path] for te in tevents]
-        #print(path)
         for depth, name in enumerate(path, 1):
             if name not in current_tree:
                 current_tree[name] = OrderedDict()
@@ -140,7 +138,8 @@ def traces_to_display(traces, trace_events, show_events=False, paths=None):
                                 sum([e.cpu_time_total for e in events if e.name == event.name]),
                                 sum([e.cuda_time_total for e in events if e.name == event.name]),
                                 len([e for e in events if e.name == event.name]),
-                                str(module)
+                                str(module),
+                                event.input_shapes
                             )._asdict()
                         }
                 else:
@@ -149,7 +148,8 @@ def traces_to_display(traces, trace_events, show_events=False, paths=None):
                         sum([e.cpu_time_total for e in events]),
                         sum([e.cuda_time_total for e in events]),
                         len(trace_events[path]),
-                        str(module)
+                        str(module),
+                        [e.input_shapes for e in events][0]
                     )._asdict()
             current_tree = current_tree[name]
     return tree