def ops2dag(merged_ops): """Load ops dict into dag.""" dag = DAG() outs = { op['outputs'].name: op for name, op in merged_ops.items() if op['outputs'] is not None } for name, node in merged_ops.items(): inps = node['inputs'] pre_node_name = 'root' dag.add_node_if_not_exists(name) if inps is not None: for inp in inps: pre_node = outs.get(inp.name) if pre_node is not None: pre_node_name = pre_node.op_name dag.add_edge(pre_node_name, name) else: dag.add_edge(pre_node_name, name) return dag
def ops2dag(merged_ops): """Load ops dict into dag.""" dag = DAG() dot = DagGraphVisual() dot.node(name='root', label='root') outs = {op['outputs'].name: op for name, op in merged_ops.items() if op['outputs'] is not None} outs = {k.replace('Conv2D:0', 'BiasAdd:0'): v for k, v in outs.items()} for name, node in merged_ops.items(): inps = node['inputs'] pre_node_name = 'root' dag.add_node_if_not_exists(name) dot.node(name=name, label=name) if inps is not None: for inp in inps: pre_node = outs.get(inp.name) if pre_node is not None: pre_node_name = pre_node.op_name dag.add_edge(pre_node_name, name) dot.edge(pre_node_name, name) else: dag.add_edge(pre_node_name, name) dot.edge(pre_node_name, name) dot.show() return dag