diff options
| author | James Steven Supancic III <arrummzen@gmail.com> | 2015-02-12 13:57:25 -0800 |
|---|---|---|
| committer | Evan Shelhamer <shelhamer@imaginarynumber.net> | 2015-02-16 11:39:14 -0800 |
| commit | 1344d1bc51f781b593a52fa71ecd59a23dc031fe (patch) | |
| tree | fe002f075c04f428f7d71929354e340f7571df91 | |
| parent | c09de3588c0a6f7376cb78cf92ba67fd9d97ba34 (diff) | |
| download | caffeonacl-1344d1bc51f781b593a52fa71ecd59a23dc031fe.tar.gz caffeonacl-1344d1bc51f781b593a52fa71ecd59a23dc031fe.tar.bz2 caffeonacl-1344d1bc51f781b593a52fa71ecd59a23dc031fe.zip | |
Fix Draw Net Problem #1709
Introduced by Layer type is a string #1694
| -rw-r--r-- | python/caffe/draw.py | 12 |
1 files changed, 2 insertions, 10 deletions
diff --git a/python/caffe/draw.py b/python/caffe/draw.py index d95c193c..1871e7a0 100644 --- a/python/caffe/draw.py +++ b/python/caffe/draw.py @@ -17,13 +17,6 @@ NEURON_LAYER_STYLE = {'shape': 'record', 'fillcolor': '#90EE90', '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 @@ -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]): |