Removes unused nodes from a graph file.
(input_graph, input_binary, output_graph,
output_binary, input_node_names, output_node_names,
placeholder_type_enum)
| 90 | |
| 91 | |
| 92 | def strip_unused_from_files(input_graph, input_binary, output_graph, |
| 93 | output_binary, input_node_names, output_node_names, |
| 94 | placeholder_type_enum): |
| 95 | """Removes unused nodes from a graph file.""" |
| 96 | |
| 97 | if not gfile.Exists(input_graph): |
| 98 | print("Input graph file '" + input_graph + "' does not exist!") |
| 99 | return -1 |
| 100 | |
| 101 | if not output_node_names: |
| 102 | print("You need to supply the name of a node to --output_node_names.") |
| 103 | return -1 |
| 104 | |
| 105 | input_graph_def = graph_pb2.GraphDef() |
| 106 | mode = "rb" if input_binary else "r" |
| 107 | with gfile.GFile(input_graph, mode) as f: |
| 108 | if input_binary: |
| 109 | input_graph_def.ParseFromString(f.read()) |
| 110 | else: |
| 111 | text_format.Merge(f.read(), input_graph_def) |
| 112 | |
| 113 | output_graph_def = strip_unused(input_graph_def, |
| 114 | input_node_names.split(","), |
| 115 | output_node_names.split(","), |
| 116 | placeholder_type_enum) |
| 117 | |
| 118 | if output_binary: |
| 119 | with gfile.GFile(output_graph, "wb") as f: |
| 120 | f.write(output_graph_def.SerializeToString()) |
| 121 | else: |
| 122 | with gfile.GFile(output_graph, "w") as f: |
| 123 | f.write(text_format.MessageToString(output_graph_def)) |
| 124 | print("%d ops in the final graph." % len(output_graph_def.node)) |
nothing calls this directly
no test coverage detected