/usr/local/lib64/python3.6/site-packages/caffe2/python
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/net_drawer.py (14264B)
## @package net_drawer
# Module caffe2.python.net_drawer
import argparse
import json
import logging
from collections import defaultdict
from caffe2.python import utils
from future.utils import viewitems
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
try:
import pydot
except ImportError:
logger.info(
'Cannot import pydot, which is required for drawing a network. This '
'can usually be installed in python with "pip install pydot". Also, '
'pydot requires graphviz to convert dot files to pdf: in ubuntu, this '
'can usually be installed with "sudo apt-get install graphviz".'
)
print(
'net_drawer will not run correctly. Please install the correct '
'dependencies.'
)
pydot = None
from caffe2.proto import caffe2_pb2
OP_STYLE = {
'shape': 'box',
'color': '#0F9D58',
'style': 'filled',
'fontcolor': '#FFFFFF'
}
BLOB_STYLE = {'shape': 'octagon'}
def _rectify_operator_and_name(operators_or_net, name):
"""Gets the operators and name for the pydot graph."""
if isinstance(operators_or_net, caffe2_pb2.NetDef):
operators = operators_or_net.op
if name is None:
name = operators_or_net.name
elif hasattr(operators_or_net, 'Proto'):
net = operators_or_net.Proto()
if not isinstance(net, caffe2_pb2.NetDef):
raise RuntimeError(
"Expecting NetDef, but got {}".format(type(net)))
operators = net.op
if name is None:
name = net.name
else:
operators = operators_or_net
if name is None:
name = "unnamed"
return operators, name
def _escape_label(name):
# json.dumps is poor man's escaping
return json.dumps(name)
def GetOpNodeProducer(append_output, **kwargs):
def ReallyGetOpNode(op, op_id):
if op.name:
node_name = '%s/%s (op#%d)' % (op.name, op.type, op_id)
else:
node_name = '%s (op#%d)' % (op.type, op_id)
if append_output:
for output_name in op.output:
node_name += '\n' + output_name
return pydot.Node(node_name, **kwargs)
return ReallyGetOpNode
def GetBlobNodeProducer(**kwargs):
def ReallyGetBlobNode(node_name, label):
return pydot.Node(node_name, label=label, **kwargs)
return ReallyGetBlobNode
def GetPydotGraph(
operators_or_net,
name=None,
rankdir='LR',
op_node_producer=None,
blob_node_producer=None
):
if op_node_producer is None:
op_node_producer = GetOpNodeProducer(False, **OP_STYLE)
if blob_node_producer is None:
blob_node_producer = GetBlobNodeProducer(**BLOB_STYLE)
operators, name = _rectify_operator_and_name(operators_or_net, name)
graph = pydot.Dot(name, rankdir=rankdir)
pydot_nodes = {}
pydot_node_counts = defaultdict(int)
for op_id, op in enumerate(operators):
op_node = op_node_producer(op, op_id)
graph.add_node(op_node)
# print 'Op: %s' % op.name
# print 'inputs: %s' % str(op.input)
# print 'outputs: %s' % str(op.output)
for input_name in op.input:
if input_name not in pydot_nodes:
input_node = blob_node_producer(
_escape_label(
input_name + str(pydot_node_counts[input_name])),
label=_escape_label(input_name),
)
pydot_nodes[input_name] = input_node
else:
input_node = pydot_nodes[input_name]
graph.add_node(input_node)
graph.add_edge(pydot.Edge(input_node, op_node))
for output_name in op.output:
if output_name in pydot_nodes:
# we are overwriting an existing blob. need to update the count.
pydot_node_counts[output_name] += 1
output_node = blob_node_producer(
_escape_label(
output_name + str(pydot_node_counts[output_name])),
label=_escape_label(output_name),
)
pydot_nodes[output_name] = output_node
graph.add_node(output_node)
graph.add_edge(pydot.Edge(op_node, output_node))
return graph
def GetPydotGraphMinimal(
operators_or_net,
name=None,
rankdir='LR',
minimal_dependency=False,
op_node_producer=None,
):
"""Different from GetPydotGraph, hide all blob nodes and only show op nodes.
If minimal_dependency is set as well, for each op, we will only draw the
edges to the minimal necessary ancestors. For example, if op c depends on
op a and b, and op b depends on a, then only the edge b->c will be drawn
because a->c will be implied.
"""
if op_node_producer is None:
op_node_producer = GetOpNodeProducer(False, **OP_STYLE)
operators, name = _rectify_operator_and_name(operators_or_net, name)
graph = pydot.Dot(name, rankdir=rankdir)
# blob_parents maps each blob name to its generating op.
blob_parents = {}
# op_ancestry records the ancestors of each op.
op_ancestry = defaultdict(set)
for op_id, op in enumerate(operators):
op_node = op_node_producer(op, op_id)
graph.add_node(op_node)
# Get parents, and set up op ancestry.
parents = [
blob_parents[input_name] for input_name in op.input
if input_name in blob_parents
]
op_ancestry[op_node].update(parents)
for node in parents:
op_ancestry[op_node].update(op_ancestry[node])
if minimal_dependency:
# only add nodes that do not have transitive ancestry
for node in parents:
if all(
[node not in op_ancestry[other_node]
for other_node in parents]
):
graph.add_edge(pydot.Edge(node, op_node))
else:
# Add all parents to the graph.
for node in parents:
graph.add_edge(pydot.Edge(node, op_node))
# Update blob_parents to reflect that this op created the blobs.
for output_name in op.output:
blob_parents[output_name] = op_node
return graph
def GetOperatorMapForPlan(plan_def):
operator_map = {}
for net_id, net in enumerate(plan_def.network):
if net.HasField('name'):
operator_map[plan_def.name + "_" + net.name] = net.op
else:
operator_map[plan_def.name + "_network_%d" % net_id] = net.op
return operator_map
def _draw_nets(nets, g):
nodes = []
for i, net in enumerate(nets):
nodes.append(pydot.Node(_escape_label(net)))
g.add_node(nodes[-1])
if i > 0:
g.add_edge(pydot.Edge(nodes[-2], nodes[-1]))
return nodes
def _draw_steps(steps, g, skip_step_edges=False): # noqa
kMaxParallelSteps = 3
def get_label():
label = [step.name + '\n']
if step.report_net:
label.append('Reporter: {}'.format(step.report_net))
if step.should_stop_blob:
label.append('Stopper: {}'.format(step.should_stop_blob))
if step.concurrent_substeps:
label.append('Concurrent')
if step.only_once:
label.append('Once')
return '\n'.join(label)
def substep_edge(start, end):
return pydot.Edge(start, end, arrowhead='dot', style='dashed')
nodes = []
for i, step in enumerate(steps):
parallel = step.concurrent_substeps
nodes.append(pydot.Node(_escape_label(get_label()), **OP_STYLE))
g.add_node(nodes[-1])
if i > 0 and not skip_step_edges:
g.add_edge(pydot.Edge(nodes[-2], nodes[-1]))
if step.network:
sub_nodes = _draw_nets(step.network, g)
elif step.substep:
if parallel:
sub_nodes = _draw_steps(
step.substep[:kMaxParallelSteps], g, skip_step_edges=True)
else:
sub_nodes = _draw_steps(step.substep, g)
else:
raise ValueError('invalid step')
if parallel:
for sn in sub_nodes:
g.add_edge(substep_edge(nodes[-1], sn))
if len(step.substep) > kMaxParallelSteps:
ellipsis = pydot.Node('{} more steps'.format(
len(step.substep) - kMaxParallelSteps), **OP_STYLE)
g.add_node(ellipsis)
g.add_edge(substep_edge(nodes[-1], ellipsis))
else:
g.add_edge(substep_edge(nodes[-1], sub_nodes[0]))
return nodes
def GetPlanGraph(plan_def, name=None, rankdir='TB'):
graph = pydot.Dot(name, rankdir=rankdir)
_draw_steps(plan_def.execution_step, graph)
return graph
def GetGraphInJson(operators_or_net, output_filepath):
operators, _ = _rectify_operator_and_name(operators_or_net, None)
blob_strid_to_node_id = {}
node_name_counts = defaultdict(int)
nodes = []
edges = []
for op_id, op in enumerate(operators):
op_label = op.name + '/' + op.type if op.name else op.type
op_node_id = len(nodes)
nodes.append({
'id': op_node_id,
'label': op_label,
'op_id': op_id,
'type': 'op'
})
for input_name in op.input:
strid = _escape_label(
input_name + str(node_name_counts[input_name]))
if strid not in blob_strid_to_node_id:
input_node = {
'id': len(nodes),
'label': input_name,
'type': 'blob'
}
blob_strid_to_node_id[strid] = len(nodes)
nodes.append(input_node)
else:
input_node = nodes[blob_strid_to_node_id[strid]]
edges.append({
'source': blob_strid_to_node_id[strid],
'target': op_node_id
})
for output_name in op.output:
strid = _escape_label(
output_name + str(node_name_counts[output_name]))
if strid in blob_strid_to_node_id:
# we are overwriting an existing blob. need to update the count.
node_name_counts[output_name] += 1
strid = _escape_label(
output_name + str(node_name_counts[output_name]))
if strid not in blob_strid_to_node_id:
output_node = {
'id': len(nodes),
'label': output_name,
'type': 'blob'
}
blob_strid_to_node_id[strid] = len(nodes)
nodes.append(output_node)
edges.append({
'source': op_node_id,
'target': blob_strid_to_node_id[strid]
})
with open(output_filepath, 'w') as f:
json.dump({'nodes': nodes, 'edges': edges}, f)
# A dummy minimal PNG image used by GetGraphPngSafe as a
# placeholder when rendering fail to run.
_DummyPngImage = (
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00'
b'\x01\x01\x00\x00\x00\x007n\xf9$\x00\x00\x00\nIDATx\x9cc`\x00\x00'
b'\x00\x02\x00\x01H\xaf\xa4q\x00\x00\x00\x00IEND\xaeB`\x82')
def GetGraphPngSafe(func, *args, **kwargs):
"""
Invokes `func` (e.g. GetPydotGraph) with args. If anything fails - returns
and empty image instead of throwing Exception
"""
try:
graph = func(*args, **kwargs)
if not isinstance(graph, pydot.Dot):
raise ValueError("func is expected to return pydot.Dot")
return graph.create_png()
except Exception as e:
logger.error("Failed to draw graph: {}".format(e))
return _DummyPngImage
def main():
parser = argparse.ArgumentParser(description="Caffe2 net drawer.")
parser.add_argument(
"--input",
type=str, required=True,
help="The input protobuf file."
)
parser.add_argument(
"--output_prefix",
type=str, default="",
help="The prefix to be added to the output filename."
)
parser.add_argument(
"--minimal", action="store_true",
help="If set, produce a minimal visualization."
)
parser.add_argument(
"--minimal_dependency", action="store_true",
help="If set, only draw minimal dependency."
)
parser.add_argument(
"--append_output", action="store_true",
help="If set, append the output blobs to the operator names.")
parser.add_argument(
"--rankdir", type=str, default="LR",
help="The rank direction of the pydot graph."
)
args = parser.parse_args()
with open(args.input, 'r') as fid:
content = fid.read()
graphs = utils.GetContentFromProtoString(
content, {
caffe2_pb2.PlanDef: lambda x: GetOperatorMapForPlan(x),
caffe2_pb2.NetDef: lambda x: {x.name: x.op},
}
)
for key, operators in viewitems(graphs):
if args.minimal:
graph = GetPydotGraphMinimal(
operators,
name=key,
rankdir=args.rankdir,
node_producer=GetOpNodeProducer(args.append_output, **OP_STYLE),
minimal_dependency=args.minimal_dependency)
else:
graph = GetPydotGraph(
operators,
name=key,
rankdir=args.rankdir,
node_producer=GetOpNodeProducer(args.append_output, **OP_STYLE))
filename = args.output_prefix + graph.get_name() + '.dot'
graph.write(filename, format='raw')
pdf_filename = filename[:-3] + 'pdf'
try:
graph.write_pdf(pdf_filename)
except Exception:
print(
'Error when writing out the pdf file. Pydot requires graphviz '
'to convert dot files to pdf, and you may not have installed '
'graphviz. On ubuntu this can usually be installed with "sudo '
'apt-get install graphviz". We have generated the .dot file '
'but will not be able to generate pdf file for now.'
)
if __name__ == '__main__':
main()