Showing the changes to tools/python/convert.py (1 of 2 files). Show all
Unified · Split
| @@ -27,6 +27,7 @@ def get_arguments(): · 26 unmodified lines | ||
| 27 | 27 | parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') |
| 28 | 28 | parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model') |
| 29 | 29 | parser.add_argument('infile', help='path to the deep learning model with weights') |
| 30 | + parser.add_argument('--dump4tb', type=str, default='no', help='dump file for visualization in tensorboard') | |
| 30 | 31 | |
| 31 | 32 | return parser.parse_args() |
| 32 | 33 | |
| @@ -44,9 +45,12 @@ def main(): · 11 unmodified lines | ||
| 44 | 45 | basefile = os.path.split(args.infile)[1] |
| 45 | 46 | basefile = os.path.splitext(basefile)[0] |
| 46 | 47 | outfile = os.path.join(args.outdir, basefile) + '.model' |
| 48 | + dump4tb = False | |
| 49 | + if args.dump4tb.lower() in ('yes', 'true', 't', 'y', '1'): | |
| 50 | + dump4tb = True | |
| 47 | 51 | |
| 48 | 52 | if args.infmt == 'tensorflow': |
| 49 | − convert_from_tensorflow(args.infile, outfile) | |
| 53 | + convert_from_tensorflow(args.infile, outfile, dump4tb) | |
| 50 | 54 | |
| 51 | 55 | if __name__ == '__main__': |
| 52 | 56 | main() |