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
(layer, rankdir)
| 60 | |
| 61 | |
| 62 | def get_layer_label(layer, rankdir): |
| 63 | """Define node label based on layer type. |
| 64 | |
| 65 | Parameters |
| 66 | ---------- |
| 67 | layer : ? |
| 68 | rankdir : {'LR', 'TB', 'BT'} |
| 69 | Direction of graph layout. |
| 70 | |
| 71 | Returns |
| 72 | ------- |
| 73 | string : |
| 74 | A label for the current layer |
| 75 | """ |
| 76 | |
| 77 | if rankdir in ('TB', 'BT'): |
| 78 | # If graph orientation is vertical, horizontal space is free and |
| 79 | # vertical space is not; separate words with spaces |
| 80 | separator = ' ' |
| 81 | else: |
| 82 | # If graph orientation is horizontal, vertical space is free and |
| 83 | # horizontal space is not; separate words with newlines |
| 84 | separator = '\\n' |
| 85 | |
| 86 | if layer.type == 'Convolution' or layer.type == 'Deconvolution': |
| 87 | # Outer double quotes needed or else colon characters don't parse |
| 88 | # properly |
| 89 | node_label = '"%s%s(%s)%skernel size: %d%sstride: %d%spad: %d"' %\ |
| 90 | (layer.name, |
| 91 | separator, |
| 92 | layer.type, |
| 93 | separator, |
| 94 | layer.convolution_param.kernel_size[0] if len(layer.convolution_param.kernel_size._values) else 1, |
| 95 | separator, |
| 96 | layer.convolution_param.stride[0] if len(layer.convolution_param.stride._values) else 1, |
| 97 | separator, |
| 98 | layer.convolution_param.pad[0] if len(layer.convolution_param.pad._values) else 0) |
| 99 | elif layer.type == 'Pooling': |
| 100 | pooling_types_dict = get_pooling_types_dict() |
| 101 | node_label = '"%s%s(%s %s)%skernel size: %d%sstride: %d%spad: %d"' %\ |
| 102 | (layer.name, |
| 103 | separator, |
| 104 | pooling_types_dict[layer.pooling_param.pool], |
| 105 | layer.type, |
| 106 | separator, |
| 107 | layer.pooling_param.kernel_size, |
| 108 | separator, |
| 109 | layer.pooling_param.stride, |
| 110 | separator, |
| 111 | layer.pooling_param.pad) |
| 112 | else: |
| 113 | node_label = '"%s%s(%s)"' % (layer.name, separator, layer.type) |
| 114 | return node_label |
| 115 | |
| 116 | |
| 117 | def choose_color_by_layertype(layertype): |
no test coverage detected