Define edge label based on layer type.
(layer)
| 44 | |
| 45 | |
| 46 | def get_edge_label(layer): |
| 47 | """Define edge label based on layer type. |
| 48 | """ |
| 49 | |
| 50 | if layer.type == 'Data': |
| 51 | edge_label = 'Batch ' + str(layer.data_param.batch_size) |
| 52 | elif layer.type == 'Convolution' or layer.type == 'Deconvolution': |
| 53 | edge_label = str(layer.convolution_param.num_output) |
| 54 | elif layer.type == 'InnerProduct': |
| 55 | edge_label = str(layer.inner_product_param.num_output) |
| 56 | else: |
| 57 | edge_label = '""' |
| 58 | |
| 59 | return edge_label |
| 60 | |
| 61 | |
| 62 | def get_layer_label(layer, rankdir): |