| |
| """ |
| Draw a graph of the net architecture. |
| """ |
| from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter |
| from google.protobuf import text_format |
|
|
| import caffe |
| import caffe.draw |
| from caffe.proto import caffe_pb2 |
|
|
|
|
| def parse_args(): |
| """Parse input arguments |
| """ |
|
|
| parser = ArgumentParser(description=__doc__, |
| formatter_class=ArgumentDefaultsHelpFormatter) |
|
|
| parser.add_argument('input_net_proto_file', |
| help='Input network prototxt file') |
| parser.add_argument('output_image_file', |
| help='Output image file') |
| parser.add_argument('--rankdir', |
| help=('One of TB (top-bottom, i.e., vertical), ' |
| 'RL (right-left, i.e., horizontal), or another ' |
| 'valid dot option; see ' |
| 'http://www.graphviz.org/doc/info/' |
| 'attrs.html#k:rankdir'), |
| default='LR') |
| parser.add_argument('--phase', |
| help=('Which network phase to draw: can be TRAIN, ' |
| 'TEST, or ALL. If ALL, then all layers are drawn ' |
| 'regardless of phase.'), |
| default="ALL") |
| parser.add_argument('--display_lrm', action='store_true', |
| help=('Use this flag to visualize the learning rate ' |
| 'multiplier, when non-zero, for the learning ' |
| 'layers (Convolution, Deconvolution, InnerProduct).')) |
|
|
| args = parser.parse_args() |
| return args |
|
|
|
|
| def main(): |
| args = parse_args() |
| net = caffe_pb2.NetParameter() |
| text_format.Merge(open(args.input_net_proto_file).read(), net) |
| print('Drawing net to %s' % args.output_image_file) |
| phase=None; |
| if args.phase == "TRAIN": |
| phase = caffe.TRAIN |
| elif args.phase == "TEST": |
| phase = caffe.TEST |
| elif args.phase != "ALL": |
| raise ValueError("Unknown phase: " + args.phase) |
| caffe.draw.draw_net_to_file(net, args.output_image_file, args.rankdir, |
| phase, args.display_lrm) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
