From c55cb6d86626aea2429c5da67720e86b056127e9 Mon Sep 17 00:00:00 2001
From: Hideaki Masuda <masuda@leapmind.io>
Date: Mon, 23 Dec 2019 11:43:00 +0900
Subject: [PATCH] use named values for node parameters

---
 lmnet/executor/profile_model.py | 56 ++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 25 deletions(-)

diff --git a/lmnet/executor/profile_model.py b/lmnet/executor/profile_model.py
index ca3fcc40a..6695c7f33 100644
--- a/lmnet/executor/profile_model.py
+++ b/lmnet/executor/profile_model.py
@@ -176,41 +176,47 @@ def helper(node, level):
                           node.name.split("/")[-1]
         bits = bit if is_quant_kernel else 32
         node_name = "total" if level == 0 else node.name
+        node_params = node.total_parameters
+        node_size = node_params * 32
+        node_quant_size = (node_params * bit) if is_quant_kernel else node_size
+        node_children = []
+        # Add node info to result list
         res.append(
-            [level, node_name, node.total_parameters, node.total_parameters * 32, node.total_parameters * bits])
+            [level, node_name, node_params, node_size, node_quant_size])
         idx = len(res) - 1
-
-        node_param_dict = {
-            'name': node_name,
-            'parameters': node.total_parameters,
-            'size': node.total_parameters * 32,
-            'quant_size': node.total_parameters * bits,
-            'children': [],
-        }
+        # Get children node info
+        sumqs = 0
+        for c in node.children:
+            children = helper(c, level + 1)
+            sumqs += children["quant_size"]
+            node_children.append(children)
+        node_quant_size = sumqs or node_quant_size
+        # Update node_quant_size of result list
+        res[idx][-1] = node_quant_size
+        # Create node_dict
         if node_name == "total":
             node_param_dict = {
-                'total_parameters': node.total_parameters,
-                'total_size': node.total_parameters * 32,
+                'total_parameters': node_params,
+                'total_size': node_size,
                 'quant_bit': bit,
-                'total_quant_size': None,
-                'children': [],
+                'total_quant_size': node_quant_size,
+                'children': node_children,
             }
-
-        sumsq = 0
-        for c in node.children:
-            size_quat, children = helper(c, level + 1)
-            sumsq += size_quat
-            node_param_dict["children"].append(children)
-        res[idx][-1] = sumsq or res[idx][-1]
-        if node_name == "total":
-            node_param_dict["total_quant_size"] = res[idx][-1]
         else:
-            node_param_dict["quant_size"] = res[idx][-1]
+            node_param_dict = {
+                'name': node_name,
+                'parameters': node_params,
+                'size': node_size,
+                'quant_size': node_quant_size,
+                'children': node_children,
+            }
+        # Add is_quant_kernel flag to leaf node
         if len(node.children) == 0:
             node_param_dict["is_quant_kernel"] = is_quant_kernel
-        return res[idx][-1], node_param_dict
 
-    _, node_param_dict = helper(prof, 0)
+        return node_param_dict
+
+    node_param_dict = helper(prof, 0)
     for elem in res:
         elem[3] = round(elem[3] / 8 / 1024 ** 2, 5)
         elem[4] = round(elem[4] / 8 / 1024 ** 2, 5)