()
| 39 | |
| 40 | |
| 41 | def main(): |
| 42 | args = parse_args() |
| 43 | net = caffe_pb2.NetParameter() |
| 44 | text_format.Merge(open(args.input_net_proto_file).read(), net) |
| 45 | print('Drawing net to %s' % args.output_image_file) |
| 46 | phase=None; |
| 47 | if args.phase == "TRAIN": |
| 48 | phase = caffe.TRAIN |
| 49 | elif args.phase == "TEST": |
| 50 | phase = caffe.TEST |
| 51 | elif args.phase != "ALL": |
| 52 | raise ValueError("Unknown phase: " + args.phase) |
| 53 | caffe.draw.draw_net_to_file(net, args.output_image_file, args.rankdir, |
| 54 | phase) |
| 55 | |
| 56 | |
| 57 | if __name__ == '__main__': |