Skip to content

Commit

Permalink
python/draw_net.py and python/caffe/draw.py: Simplified code; added m…
Browse files Browse the repository at this point in the history 
…ore docstrings; adjusted code according to PEP8
  • Loading branch information
@MartinThoma
MartinThoma committed May 27, 2015
1 parent e8d93cb commit adf3189
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 67 deletions.
163 changes: 102 additions & 61 deletions python/caffe/draw.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
"""
Caffe network visualization: draw the NetParameter protobuffer.
NOTE: this requires pydot>=1.0.2, which is not included in requirements.txt
since it requires graphviz and other prerequisites outside the scope of the
Caffe.
.. note::
This requires pydot>=1.0.2, which is not included in requirements.txt since
it requires graphviz and other prerequisites outside the scope of the
Caffe.
"""

from caffe.proto import caffe_pb2
from google.protobuf import text_format
import pydot

# Internal layer and blob styles.
Expand All @@ -32,24 +34,35 @@ def get_pooling_types_dict():
return d


def determine_edge_label_by_layertype(layer, layertype):
"""Define edge label based on layer type
def get_edge_label(layer):
"""Define edge label based on layer type.
"""

if layertype == 'Data':
if layer.type == 'Data':
edge_label = 'Batch ' + str(layer.data_param.batch_size)
elif layertype == 'Convolution':
elif layer.type == 'Convolution':
edge_label = str(layer.convolution_param.num_output)
elif layertype == 'InnerProduct':
elif layer.type == 'InnerProduct':
edge_label = str(layer.inner_product_param.num_output)
else:
edge_label = '""'

return edge_label


def determine_node_label_by_layertype(layer, layertype, rankdir):
"""Define node label based on layer type
def get_layer_label(layer, rankdir):
"""Define node label based on layer type.
Parameters
----------
layer : ?
rankdir : {'LR', 'TB', 'BT'}
Direction of graph layout.
Returns
-------
string :
A label for the current layer
"""

if rankdir in ('TB', 'BT'):
Expand All @@ -61,39 +74,39 @@ def determine_node_label_by_layertype(layer, layertype, rankdir):
# horizontal space is not; separate words with newlines
separator = '\n'

if layertype == 'Convolution':
if layer.type == 'Convolution':
# Outer double quotes needed or else colon characters don't parse
# properly
node_label = '"%s%s(%s)%skernel size: %d%sstride: %d%spad: %d"' %\
(layer.name,
separator,
layertype,
layer.type,
separator,
layer.convolution_param.kernel_size,
separator,
layer.convolution_param.stride,
separator,
layer.convolution_param.pad)
elif layertype == 'Pooling':
elif layer.type == 'Pooling':
pooling_types_dict = get_pooling_types_dict()
node_label = '"%s%s(%s %s)%skernel size: %d%sstride: %d%spad: %d"' %\
(layer.name,
separator,
pooling_types_dict[layer.pooling_param.pool],
layertype,
layer.type,
separator,
layer.pooling_param.kernel_size,
separator,
layer.pooling_param.stride,
separator,
layer.pooling_param.pad)
else:
node_label = '"%s%s(%s)"' % (layer.name, separator, layertype)
node_label = '"%s%s(%s)"' % (layer.name, separator, layer.type)
return node_label


def choose_color_by_layertype(layertype):
"""Define colors for nodes based on the layer type
"""Define colors for nodes based on the layer type.
"""
color = '#6495ED' # Default
if layertype == 'Convolution':
Expand All @@ -106,48 +119,62 @@ def choose_color_by_layertype(layertype):


def get_pydot_graph(caffe_net, rankdir, label_edges=True):
pydot_graph = pydot.Dot(caffe_net.name, graph_type='digraph', rankdir=rankdir)
pydot_nodes = {}
pydot_edges = []
for layer in caffe_net.layer:
name = layer.name
layertype = layer.type
node_label = determine_node_label_by_layertype(layer, layertype, rankdir)
if (len(layer.bottom) == 1 and len(layer.top) == 1 and
layer.bottom[0] == layer.top[0]):
# We have an in-place neuron layer.
pydot_nodes[name + '_' + layertype] = pydot.Node(
node_label, **NEURON_LAYER_STYLE)
else:
layer_style = LAYER_STYLE_DEFAULT
layer_style['fillcolor'] = choose_color_by_layertype(layertype)
pydot_nodes[name + '_' + layertype] = pydot.Node(
node_label, **layer_style)
for bottom_blob in layer.bottom:
pydot_nodes[bottom_blob + '_blob'] = pydot.Node(
'%s' % (bottom_blob), **BLOB_STYLE)
edge_label = '""'
pydot_edges.append({'src': bottom_blob + '_blob',
'dst': name + '_' + layertype,
'label': edge_label})
for top_blob in layer.top:
pydot_nodes[top_blob + '_blob'] = pydot.Node(
'%s' % (top_blob))
if label_edges:
edge_label = determine_edge_label_by_layertype(layer, layertype)
else:
edge_label = '""'
pydot_edges.append({'src': name + '_' + layertype,
'dst': top_blob + '_blob',
'label': edge_label})
# Now, add the nodes and edges to the graph.
for node in pydot_nodes.values():
pydot_graph.add_node(node)
for edge in pydot_edges:
pydot_graph.add_edge(
pydot.Edge(pydot_nodes[edge['src']], pydot_nodes[edge['dst']],
label=edge['label']))
return pydot_graph
"""Create a data structure which represents the `caffe_net`.
Parameters
----------
caffe_net : object
rankdir : {'LR', 'TB', 'BT'}
Direction of graph layout.
label_edges : boolean, optional
Label the edges (default is True).
Returns
-------
pydot graph object
"""
pydot_graph = pydot.Dot(caffe_net.name,
graph_type='digraph',
rankdir=rankdir)
pydot_nodes = {}
pydot_edges = []
for layer in caffe_net.layer:
node_label = get_layer_label(layer, rankdir)
node_name = "%s_%s" % (layer.name, layer.type)
if (len(layer.bottom) == 1 and len(layer.top) == 1 and
layer.bottom[0] == layer.top[0]):
# We have an in-place neuron layer.
pydot_nodes[node_name] = pydot.Node(node_label,
**NEURON_LAYER_STYLE)
else:
layer_style = LAYER_STYLE_DEFAULT
layer_style['fillcolor'] = choose_color_by_layertype(layer.type)
pydot_nodes[node_name] = pydot.Node(node_label, **layer_style)
for bottom_blob in layer.bottom:
pydot_nodes[bottom_blob + '_blob'] = pydot.Node('%s' % bottom_blob,
**BLOB_STYLE)
edge_label = '""'
pydot_edges.append({'src': bottom_blob + '_blob',
'dst': node_name,
'label': edge_label})
for top_blob in layer.top:
pydot_nodes[top_blob + '_blob'] = pydot.Node('%s' % (top_blob))
if label_edges:
edge_label = get_edge_label(layer)
else:
edge_label = '""'
pydot_edges.append({'src': node_name,
'dst': top_blob + '_blob',
'label': edge_label})
# Now, add the nodes and edges to the graph.
for node in pydot_nodes.values():
pydot_graph.add_node(node)
for edge in pydot_edges:
pydot_graph.add_edge(
pydot.Edge(pydot_nodes[edge['src']],
pydot_nodes[edge['dst']],
label=edge['label']))
return pydot_graph


def draw_net(caffe_net, rankdir, ext='png'):
Expand All @@ -156,8 +183,14 @@ def draw_net(caffe_net, rankdir, ext='png'):
Parameters
----------
caffe_net: a caffe.proto.caffe_pb2.NetParameter protocol buffer.
ext: the image extension. Default 'png'.
caffe_net : a caffe.proto.caffe_pb2.NetParameter protocol buffer.
ext : string, optional
The image extension (the default is 'png').
Returns
-------
string :
Postscript representation of the graph.
"""
return get_pydot_graph(caffe_net, rankdir).create(format=ext)

Expand All @@ -166,6 +199,14 @@ def draw_net_to_file(caffe_net, filename, rankdir='LR'):
"""Draws a caffe net, and saves it to file using the format given as the
file extension. Use '.raw' to output raw text that you can manually feed
to graphviz to draw graphs.
Parameters
----------
caffe_net : a caffe.proto.caffe_pb2.NetParameter protocol buffer.
filename : string
The path to a file where the networks visualization will be stored.
rankdir : {'LR', 'TB', 'BT'}
Direction of graph layout.
"""
ext = filename[filename.rfind('.')+1:]
with open(filename, 'wb') as fid:
Expand Down
13 changes: 7 additions & 6 deletions python/draw_net.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
"""
Draw a graph of the net architecture.
"""
import argparse
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from google.protobuf import text_format

import caffe
Expand All @@ -14,18 +14,19 @@ def parse_args():
"""Parse input arguments
"""

parser = argparse.ArgumentParser(description='Draw a network graph')
parser = ArgumentParser(description=__doc__,
formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('input_net_proto_file',
help='Input network prototxt file')
parser.add_argument('output_image_file',
help='Output image file')
parser.add_argument('--rankdir',
help=('One of TB (top-bottom, i.e., vertical), '
'RL (right-left, i.e., horizontal), or another'
'valid dot option; see'
'http://www.graphviz.org/doc/info/attrs.html#k:rankdir'
'(default: LR)'),
'RL (right-left, i.e., horizontal), or another '
'valid dot option; see '
'http://www.graphviz.org/doc/info/'
'attrs.html#k:rankdir'),
default='LR')

args = parser.parse_args()
Expand Down

0 comments on commit adf3189

Please sign in to comment.