Fuse subgraph between input_nodes and output_nodes into a single custom op. Args: graph_def: A graph_pb2.GraphDef proto. input_nodes: input nodes to the subgraph to be fused. output_nodes: output nodes to the subgraph to be fused. output_dtypes: A list of output datatypes for the
(graph_def, input_nodes, output_nodes, output_dtypes,
output_quantized, op_name, op_type)
| 35 | |
| 36 | |
| 37 | def fuse_op(graph_def, input_nodes, output_nodes, output_dtypes, |
| 38 | output_quantized, op_name, op_type): |
| 39 | """Fuse subgraph between input_nodes and output_nodes into a single custom op. |
| 40 | |
| 41 | Args: |
| 42 | graph_def: A graph_pb2.GraphDef proto. |
| 43 | input_nodes: input nodes to the subgraph to be fused. |
| 44 | output_nodes: output nodes to the subgraph to be fused. |
| 45 | output_dtypes: A list of output datatypes for the custom op |
| 46 | output_quantized: A boolean flag that indicates if output is quantized |
| 47 | op_name: fused op name. |
| 48 | op_type: fused op type. |
| 49 | Returns: |
| 50 | The GraphDef of the new graph. |
| 51 | |
| 52 | Raises: |
| 53 | TypeError: If 'graph_def' is not a graph_pb2.GraphDef proto. |
| 54 | """ |
| 55 | |
| 56 | if not isinstance(graph_def, graph_pb2.GraphDef): |
| 57 | raise TypeError("graph_def must be a graph_pb2.GraphDef proto.") |
| 58 | |
| 59 | if isinstance(input_nodes, six.string_types): |
| 60 | raise TypeError("input_nodes must be a list.") |
| 61 | |
| 62 | if isinstance(output_nodes, six.string_types): |
| 63 | raise TypeError("output_nodes must be a list.") |
| 64 | |
| 65 | name_to_input_name, name_to_node, name_to_seq_num = _extract_graph_summary( |
| 66 | graph_def) |
| 67 | _assert_nodes_are_present(name_to_node, input_nodes + output_nodes) |
| 68 | |
| 69 | # Nodes upto and including input_nodes |
| 70 | reachable_by_input = _bfs_for_reachable_nodes(input_nodes, name_to_input_name) |
| 71 | # Nodes upto and including output_nodes |
| 72 | reachable_by_output = _bfs_for_reachable_nodes(output_nodes, |
| 73 | name_to_input_name) |
| 74 | |
| 75 | # Set of nodes in the list input_nodes |
| 76 | input_nodes_set = set(input_nodes) |
| 77 | |
| 78 | # Set of nodes in the list output_nodes |
| 79 | output_nodes_set = set(output_nodes) |
| 80 | |
| 81 | nodes_post_output = [] |
| 82 | for node in graph_def.node: |
| 83 | n = _node_name(node.name) |
| 84 | if n in reachable_by_output: |
| 85 | if n not in reachable_by_input and n not in output_nodes_set: |
| 86 | # n is between input and output, i.e., part of the fused op |
| 87 | next_to_visit = [n] |
| 88 | visited = set() |
| 89 | while next_to_visit: |
| 90 | cur_node = next_to_visit[0] |
| 91 | visited.add(cur_node) |
| 92 | del next_to_visit[0] |
| 93 | if cur_node in reachable_by_input and cur_node not in input_nodes_set: |
| 94 | raise TypeError("Node %s uses input %s not in input_nodes." % |
no test coverage detected