Collect the list of ops used by a graph. Does not validate that the ops are all registered. Args: graph_def: A `GraphDef` proto, as from `graph.as_graph_def()`. Returns: A list of strings, each naming an op used by the graph.
(graph_def)
| 136 | |
| 137 | |
| 138 | def ops_used_by_graph_def(graph_def): |
| 139 | """Collect the list of ops used by a graph. |
| 140 | |
| 141 | Does not validate that the ops are all registered. |
| 142 | |
| 143 | Args: |
| 144 | graph_def: A `GraphDef` proto, as from `graph.as_graph_def()`. |
| 145 | |
| 146 | Returns: |
| 147 | A list of strings, each naming an op used by the graph. |
| 148 | """ |
| 149 | # Map function names to definitions |
| 150 | name_to_function = {} |
| 151 | for fun in graph_def.library.function: |
| 152 | name_to_function[fun.signature.name] = fun |
| 153 | |
| 154 | # Collect the list of op names. Since functions can reference functions, we |
| 155 | # need a recursive traversal. |
| 156 | used_ops = set() # Includes both primitive ops and functions |
| 157 | functions_to_process = [] # A subset of used_ops |
| 158 | |
| 159 | def mark_op_as_used(op): |
| 160 | if op not in used_ops and op in name_to_function: |
| 161 | functions_to_process.append(name_to_function[op]) |
| 162 | used_ops.add(op) |
| 163 | |
| 164 | for node in graph_def.node: |
| 165 | mark_op_as_used(node.op) |
| 166 | while functions_to_process: |
| 167 | fun = functions_to_process.pop() |
| 168 | for node in fun.node_def: |
| 169 | mark_op_as_used(node.op) |
| 170 | |
| 171 | return [op for op in used_ops if op not in name_to_function] |
| 172 | |
| 173 | |
| 174 | def stripped_op_list_for_graph(graph_def): |
no test coverage detected