[InfCost] per-node norm mac/param counts, always floats for json

Author: maltanar

src/qonnx/analysis/inference_cost.py

@@ -117,6 +117,8 @@ def inference_cost_conv(model, node, discount_sparsity):
     mac_op_type_str = "op_mac_%s_%s" % (idt_name, wdt_name)
     w_mem_type_str = "mem_w_%s" % (wdt_name)
     o_mem_type_str = "mem_o_%s" % (odt_name)
+    # keep in floats to remain compatible with json serialization
+    n_macs, w_mem, o_mem = float(n_macs), float(w_mem), float(o_mem)
     ret = {mac_op_type_str: n_macs, w_mem_type_str: w_mem, o_mem_type_str: o_mem}
     return ret
 
@@ -161,6 +163,8 @@ def inference_cost_matmul(model, node, discount_sparsity):
     mac_op_type_str = "op_mac_%s_%s" % (idt_name, wdt_name)
     w_mem_type_str = "mem_w_%s" % (wdt_name)
     o_mem_type_str = "mem_o_%s" % (odt_name)
+    # keep in floats to remain compatible with json serialization
+    n_macs, w_mem, o_mem = float(n_macs), float(w_mem), float(o_mem)
     ret = {mac_op_type_str: n_macs, w_mem_type_str: w_mem, o_mem_type_str: o_mem}
     return ret
 
@@ -197,6 +201,8 @@ def inference_cost_upsample(model, node, discount_sparsity):
     mac_op_type_str = "op_mac_%s_%s" % (idt_name, idt_name)
     o_mem_type_str = "mem_o_%s" % (odt_name)
 
+    # keep in floats to remain compatible with json serialization
+    n_macs, o_mem = float(n_macs), float(o_mem)
     ret = {mac_op_type_str: n_macs, o_mem_type_str: o_mem}
     return ret
 

src/qonnx/util/inference_cost.py

@@ -99,7 +99,9 @@ def inference_cost(
     :param preprocess: If set, run preprocessing steps such as shape inference,
         datatype inference and constant folding. Strongly recommended.
     :param discount_sparsity: If set, will discount op cost of MAC ops with a
-        constant zero weight, and the mem cost of constant zero weights."""
+        constant zero weight, and the mem cost of constant zero weights.
+    :param cost_breakdown: If set, include per-node (by name) and per-node-type
+        breakdowns as part of the returned inference cost dict."""
 
     combined_results = {}
     if isinstance(model_filename_or_wrapper, ModelWrapper):
@@ -130,26 +132,19 @@ def inference_cost(
             res["total_macs"] = macs
             if "unsupported" in res:
                 res["unsupported"] = str(res["unsupported"])
-            if output_json is not None:
-                with open(output_json, "w") as f:
-                    json.dump(res, f, sort_keys=True, indent=2)
             combined_results[i] = res
-        elif i == "optype_cost":
-            per_optype_breakdown = {}
+        else:
+            per_optype_or_node_breakdown = {}
             for optype, op_res in res.items():
                 bops, macs = compute_bops_and_macs(op_res)
                 op_res = assign_mem_bits_and_elems(op_res)
                 op_res["total_bops"] = bops
                 op_res["total_macs"] = macs
-                per_optype_breakdown[optype] = op_res
-            combined_results[i] = per_optype_breakdown
-        else:
-            per_node_breakdown = {}
-            for node_name in res.keys():
-                node_res = res[node_name]
-                node_res = assign_mem_bits_and_elems(node_res)
-                per_node_breakdown[node_name] = node_res
-            combined_results[i] = per_node_breakdown
+                per_optype_or_node_breakdown[optype] = op_res
+            combined_results[i] = per_optype_or_node_breakdown
+    if output_json is not None:
+        with open(output_json, "w") as f:
+            json.dump(combined_results, f, sort_keys=True, indent=2)
     return combined_results
 
 

tests/analysis/test_inference_cost_breakdown.py

@@ -76,7 +76,9 @@ def test_inference_cost_breakdown(test_model):
     test_details = model_details[test_model]
     model = download_model(test_model, do_cleanup=True, return_modelwrapper=True)
     inf_cost = infca(model, discount_sparsity=False, cost_breakdown=True)
-    print(inf_cost.keys())
+    assert inf_cost["node_cost"]["Conv_0"]["total_macs"] == 118013952
+    assert inf_cost["node_cost"]["Conv_1"]["total_macs"] == 115605504
+    assert inf_cost["optype_cost"]["Conv"]["total_macs"] == 1813561344
     t_cost = inf_cost["total_cost"]  # total cost
     op_cost = aggregate_dict_keys(inf_cost["optype_cost"])  # cost per optype
     n_cost = aggregate_dict_keys(inf_cost["node_cost"])  # cost per node.