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). phase : {caffe_pb2.Phase.TRAIN, caffe
(caffe_net, rankdir, label_edges=True, phase=None)
| 128 | |
| 129 | |
| 130 | def get_pydot_graph(caffe_net, rankdir, label_edges=True, phase=None): |
| 131 | """Create a data structure which represents the `caffe_net`. |
| 132 | |
| 133 | Parameters |
| 134 | ---------- |
| 135 | caffe_net : object |
| 136 | rankdir : {'LR', 'TB', 'BT'} |
| 137 | Direction of graph layout. |
| 138 | label_edges : boolean, optional |
| 139 | Label the edges (default is True). |
| 140 | phase : {caffe_pb2.Phase.TRAIN, caffe_pb2.Phase.TEST, None} optional |
| 141 | Include layers from this network phase. If None, include all layers. |
| 142 | (the default is None) |
| 143 | |
| 144 | Returns |
| 145 | ------- |
| 146 | pydot graph object |
| 147 | """ |
| 148 | pydot_graph = pydot.Dot(caffe_net.name if caffe_net.name else 'Net', |
| 149 | graph_type='digraph', |
| 150 | rankdir=rankdir) |
| 151 | pydot_nodes = {} |
| 152 | pydot_edges = [] |
| 153 | for layer in caffe_net.layer: |
| 154 | if phase is not None: |
| 155 | included = False |
| 156 | if len(layer.include) == 0: |
| 157 | included = True |
| 158 | if len(layer.include) > 0 and len(layer.exclude) > 0: |
| 159 | raise ValueError('layer ' + layer.name + ' has both include ' |
| 160 | 'and exclude specified.') |
| 161 | for layer_phase in layer.include: |
| 162 | included = included or layer_phase.phase == phase |
| 163 | for layer_phase in layer.exclude: |
| 164 | included = included and not layer_phase.phase == phase |
| 165 | if not included: |
| 166 | continue |
| 167 | node_label = get_layer_label(layer, rankdir) |
| 168 | node_name = "%s_%s" % (layer.name, layer.type) |
| 169 | if (len(layer.bottom) == 1 and len(layer.top) == 1 and |
| 170 | layer.bottom[0] == layer.top[0]): |
| 171 | # We have an in-place neuron layer. |
| 172 | pydot_nodes[node_name] = pydot.Node(node_label, |
| 173 | **NEURON_LAYER_STYLE) |
| 174 | else: |
| 175 | layer_style = LAYER_STYLE_DEFAULT |
| 176 | layer_style['fillcolor'] = choose_color_by_layertype(layer.type) |
| 177 | pydot_nodes[node_name] = pydot.Node(node_label, **layer_style) |
| 178 | for bottom_blob in layer.bottom: |
| 179 | pydot_nodes[bottom_blob + '_blob'] = pydot.Node('%s' % bottom_blob, |
| 180 | **BLOB_STYLE) |
| 181 | edge_label = '""' |
| 182 | pydot_edges.append({'src': bottom_blob + '_blob', |
| 183 | 'dst': node_name, |
| 184 | 'label': edge_label}) |
| 185 | for top_blob in layer.top: |
| 186 | pydot_nodes[top_blob + '_blob'] = pydot.Node('%s' % (top_blob)) |
| 187 | if label_edges: |
no test coverage detected