Removes unused nodes from a GraphDef. Args: input_graph_def: A graph with nodes we want to prune. input_node_names: A list of the nodes we use as inputs. output_node_names: A list of the output nodes. placeholder_type_enum: The AttrValue enum for the placeholder data type, or
(input_graph_def, input_node_names, output_node_names,
placeholder_type_enum)
| 30 | |
| 31 | |
| 32 | def strip_unused(input_graph_def, input_node_names, output_node_names, |
| 33 | placeholder_type_enum): |
| 34 | """Removes unused nodes from a GraphDef. |
| 35 | |
| 36 | Args: |
| 37 | input_graph_def: A graph with nodes we want to prune. |
| 38 | input_node_names: A list of the nodes we use as inputs. |
| 39 | output_node_names: A list of the output nodes. |
| 40 | placeholder_type_enum: The AttrValue enum for the placeholder data type, or |
| 41 | a list that specifies one value per input node name. |
| 42 | |
| 43 | Returns: |
| 44 | A `GraphDef` with all unnecessary ops removed. |
| 45 | |
| 46 | Raises: |
| 47 | ValueError: If any element in `input_node_names` refers to a tensor instead |
| 48 | of an operation. |
| 49 | KeyError: If any element in `input_node_names` is not found in the graph. |
| 50 | """ |
| 51 | for name in input_node_names: |
| 52 | if ":" in name: |
| 53 | raise ValueError("Name '%s' appears to refer to a Tensor, " |
| 54 | "not a Operation." % name) |
| 55 | |
| 56 | # Here we replace the nodes we're going to override as inputs with |
| 57 | # placeholders so that any unused nodes that are inputs to them are |
| 58 | # automatically stripped out by extract_sub_graph(). |
| 59 | not_found = {name for name in input_node_names} |
| 60 | inputs_replaced_graph_def = graph_pb2.GraphDef() |
| 61 | for node in input_graph_def.node: |
| 62 | if node.name in input_node_names: |
| 63 | not_found.remove(node.name) |
| 64 | placeholder_node = node_def_pb2.NodeDef() |
| 65 | placeholder_node.op = "Placeholder" |
| 66 | placeholder_node.name = node.name |
| 67 | if isinstance(placeholder_type_enum, list): |
| 68 | input_node_index = input_node_names.index(node.name) |
| 69 | placeholder_node.attr["dtype"].CopyFrom( |
| 70 | attr_value_pb2.AttrValue(type=placeholder_type_enum[ |
| 71 | input_node_index])) |
| 72 | else: |
| 73 | placeholder_node.attr["dtype"].CopyFrom( |
| 74 | attr_value_pb2.AttrValue(type=placeholder_type_enum)) |
| 75 | if "_output_shapes" in node.attr: |
| 76 | placeholder_node.attr["_output_shapes"].CopyFrom(node.attr[ |
| 77 | "_output_shapes"]) |
| 78 | if "shape" in node.attr: |
| 79 | placeholder_node.attr["shape"].CopyFrom(node.attr["shape"]) |
| 80 | inputs_replaced_graph_def.node.extend([placeholder_node]) |
| 81 | else: |
| 82 | inputs_replaced_graph_def.node.extend([copy.deepcopy(node)]) |
| 83 | |
| 84 | if not_found: |
| 85 | raise KeyError("The following input nodes were not found: %s" % not_found) |
| 86 | |
| 87 | output_graph_def = graph_util.extract_sub_graph(inputs_replaced_graph_def, |
| 88 | output_node_names) |
| 89 | return output_graph_def |
no test coverage detected