Fixed python script for drawing nets following original issues BVLC#3698

BVLC#5477 and BVLC#5487
Bardo91 · Jun 1, 2017 · 12a5c3f · 12a5c3f
1 parent 6db0cb2
commit 12a5c3f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/python/caffe/draw.py b/python/caffe/draw.py
@@ -90,12 +90,12 @@ def get_layer_label(layer, rankdir):
                      (layer.name,
                       separator,
                       layer.type,
+                      separator,	
+                      layer.convolution_param.kernel_size[0] if len(layer.convolution_param.kernel_size) else 1,
                       separator,
-                      layer.convolution_param.kernel_size[0] if len(layer.convolution_param.kernel_size._values) else 1,
+                      layer.convolution_param.stride[0] if len(layer.convolution_param.stride) else 1,
                       separator,
-                      layer.convolution_param.stride[0] if len(layer.convolution_param.stride._values) else 1,
-                      separator,
-                      layer.convolution_param.pad[0] if len(layer.convolution_param.pad._values) else 0)
+                      layer.convolution_param.pad[0] if len(layer.convolution_param.pad) else 0)
     elif layer.type == 'Pooling':
         pooling_types_dict = get_pooling_types_dict()
         node_label = '"%s%s(%s %s)%skernel size: %d%sstride: %d%spad: %d"' %\