Skip to content

Commit

Permalink
Fix Draw Net Problem BVLC#1709
Browse files Browse the repository at this point in the history
Introduced by  Layer type is a string BVLC#1694
  • Loading branch information
jsupancic authored and shelhamer committed Feb 16, 2015
1 parent c09de35 commit 1344d1b
Showing 1 changed file with 2 additions and 10 deletions.
12 changes: 2 additions & 10 deletions python/caffe/draw.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,6 @@
'style': 'filled'}
BLOB_STYLE = {'shape': 'octagon', 'fillcolor': '#E0E0E0',
'style': 'filled'}
def get_enum_name_by_value():
desc = caffe_pb2.LayerParameter.LayerType.DESCRIPTOR
d = {}
for k,v in desc.values_by_name.items():
d[v.number] = k
return d


def get_pooling_types_dict():
"""Get dictionary mapping pooling type number to type name
Expand Down Expand Up @@ -112,10 +105,9 @@ 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 = []
d = get_enum_name_by_value()
for layer in caffe_net.layers:
for layer in caffe_net.layer:
name = layer.name
layertype = d[layer.type]
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]):
Expand Down

0 comments on commit 1344d1b

Please sign in to comment.