Substitute a custom op for the subgraph delimited by inputs and outputs.
(graph_def, inputs, outputs, op_definition)
| 76 | |
| 77 | |
| 78 | def _collapse_subgraph(graph_def, inputs, outputs, op_definition): |
| 79 | """Substitute a custom op for the subgraph delimited by inputs and outputs.""" |
| 80 | name = _uuid.uuid1().hex |
| 81 | # We need a default type, but it can be changed using 'op_definition'. |
| 82 | default_type = types_pb2.DT_FLOAT |
| 83 | new_graph = fuse_op( |
| 84 | graph_def=graph_def, |
| 85 | input_nodes=inputs, |
| 86 | output_nodes=outputs, |
| 87 | output_dtypes=[default_type for _ in outputs], |
| 88 | output_quantized=False, |
| 89 | op_name=name, |
| 90 | op_type="CustomTfLiteOp") |
| 91 | node_def = node_def_pb2.NodeDef() |
| 92 | text_format.Parse(op_definition, node_def) |
| 93 | for node in new_graph.node: |
| 94 | if node.name == name: |
| 95 | node.MergeFrom(node_def) |
| 96 | return new_graph |
| 97 | |
| 98 | |
| 99 | def main(argv): |