From e7d897a6d6da4d8d26cb69d21714545be4764786 Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Fri, 26 Oct 2018 11:53:47 +0300
Subject: [PATCH] model summary: Added tensor sizes annotation to generated
model graph PNG
---
apputils/model_summaries.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/apputils/model_summaries.py b/apputils/model_summaries.py
index 6c2cda3..1531116 100755
--- a/apputils/model_summaries.py
+++ b/apputils/model_summaries.py
@@ -489,7 +489,7 @@ def create_pydot_graph(op_nodes, data_nodes, param_nodes, edges, rankdir='TB', s
pydot_graph.add_node(pydot.Node(op_node[0], **style, label="\n".join(op_node)))
for data_node in data_nodes:
- pydot_graph.add_node(pydot.Node(data_node[0], label="\n".join(data_node)))
+ pydot_graph.add_node(pydot.Node(data_node[0], label="\n".join(data_node[1:])))
node_style = {'shape': 'oval',
'fillcolor': 'gray',
@@ -497,7 +497,7 @@ def create_pydot_graph(op_nodes, data_nodes, param_nodes, edges, rankdir='TB', s
if param_nodes is not None:
for param_node in param_nodes:
- pydot_graph.add_node(pydot.Node(param_node[0], **node_style, label="\n".join(param_node)))
+ pydot_graph.add_node(pydot.Node(param_node[0], **node_style, label="\n".join(param_node[1:])))
for edge in edges:
pydot_graph.add_edge(pydot.Edge(edge[0], edge[1]))
@@ -525,7 +525,7 @@ def create_png(sgraph, display_param_nodes=False, rankdir='TB', styles=None):
data_nodes = []
param_nodes = []
for id, param in sgraph.params.items():
- n_data = (id, str(param['shape']))
+ n_data = (id, str(distiller.volume(param['shape'])), str(param['shape']))
if data_node_has_parent(sgraph, id):
data_nodes.append(n_data)
else:
--
GitLab