Showing the changes to tools/python (2 of 2 files). Show all
Unified · Split
| @@ -27,6 +27,7 @@ def get_arguments(): · 26 unmodified lines | |||
| 27 | parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') | 27 | parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') |
| 28 | parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model') | 28 | parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model') |
| 29 | parser.add_argument('infile', help='path to the deep learning model with weights') | 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 | return parser.parse_args() | 32 | return parser.parse_args() |
| 32 | 33 | ||
| @@ -44,9 +45,12 @@ def main(): · 11 unmodified lines | |||
| 44 | basefile = os.path.split(args.infile)[1] | 45 | basefile = os.path.split(args.infile)[1] |
| 45 | basefile = os.path.splitext(basefile)[0] | 46 | basefile = os.path.splitext(basefile)[0] |
| 46 | outfile = os.path.join(args.outdir, basefile) + '.model' | 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 | if args.infmt == 'tensorflow': | 52 | if args.infmt == 'tensorflow': |
| 49 | convert_from_tensorflow(args.infile, outfile) | 53 | convert_from_tensorflow(args.infile, outfile, dump4tb) |
| 50 | 54 | ||
| 51 | if __name__ == '__main__': | 55 | if __name__ == '__main__': |
| 52 | main() | 56 | main() |
| @@ -24,10 +24,11 @@ import sys, struct · 23 unmodified lines | |||
| 24 | __all__ = ['convert_from_tensorflow'] | 24 | __all__ = ['convert_from_tensorflow'] |
| 25 | 25 | ||
| 26 | class TFConverter: | 26 | class TFConverter: |
| 27 | def __init__(self, graph_def, nodes, outfile): | 27 | def __init__(self, graph_def, nodes, outfile, dump4tb): |
| 28 | self.graph_def = graph_def | 28 | self.graph_def = graph_def |
| 29 | self.nodes = nodes | 29 | self.nodes = nodes |
| 30 | self.outfile = outfile | 30 | self.outfile = outfile |
| 31 | self.dump4tb = dump4tb | ||
| 31 | self.layer_number = 0 | 32 | self.layer_number = 0 |
| 32 | self.output_names = [] | 33 | self.output_names = [] |
| 33 | self.name_node_dict = {} | 34 | self.name_node_dict = {} |
| @@ -42,8 +43,8 @@ class TFConverter: · 8 unmodified lines | |||
| 42 | def dump_for_tensorboard(self): | 43 | def dump_for_tensorboard(self): |
| 43 | graph = tf.get_default_graph() | 44 | graph = tf.get_default_graph() |
| 44 | tf.import_graph_def(self.graph_def, name="") | 45 | tf.import_graph_def(self.graph_def, name="") |
| 45 | # tensorboard --logdir=/tmp/graph | ||
| 46 | tf.summary.FileWriter('/tmp/graph', graph) | 46 | tf.summary.FileWriter('/tmp/graph', graph) |
| 47 | print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it') | ||
| 47 | 48 | ||
| 48 | 49 | ||
| 49 | def get_conv2d_params(self, node): | 50 | def get_conv2d_params(self, node): |
| @@ -197,18 +198,18 @@ class TFConverter: · 147 unmodified lines | |||
| 197 | self.remove_identity() | 198 | self.remove_identity() |
| 198 | self.generate_edges() | 199 | self.generate_edges() |
| 199 | 200 | ||
| 200 | #check the graph with tensorboard with human eyes | 201 | if self.dump4tb: |
| 201 | #self.dump_for_tensorboard() | 202 | self.dump_for_tensorboard() |
| 202 | 203 | ||
| 203 | self.dump_to_file() | 204 | self.dump_to_file() |
| 204 | 205 | ||
| 205 | 206 | ||
| 206 | def convert_from_tensorflow(infile, outfile): | 207 | def convert_from_tensorflow(infile, outfile, dump4tb): |
| 207 | with open(infile, 'rb') as f: | 208 | with open(infile, 'rb') as f: |
| 208 | # read the file in .proto format | 209 | # read the file in .proto format |
| 209 | graph_def = tf.GraphDef() | 210 | graph_def = tf.GraphDef() |
| 210 | graph_def.ParseFromString(f.read()) | 211 | graph_def.ParseFromString(f.read()) |
| 211 | nodes = graph_def.node | 212 | nodes = graph_def.node |
| 212 | 213 | ||
| 213 | converter = TFConverter(graph_def, nodes, outfile) | 214 | converter = TFConverter(graph_def, nodes, outfile, dump4tb) |
| 214 | converter.run() | 215 | converter.run() |