Extract trainable model parameters and FLOPs for ops from a Graph. Args: graph: tf.Graph. run_meta: RunMetadata proto used to complete shape information. add_trace: Whether to add op trace information. add_trainable_var: Whether to assign tf.compat.v1.trainable_variables() op
(graph, run_meta=None, add_trace=True,
add_trainable_var=True)
| 75 | |
| 76 | |
| 77 | def _get_logged_ops(graph, run_meta=None, add_trace=True, |
| 78 | add_trainable_var=True): |
| 79 | """Extract trainable model parameters and FLOPs for ops from a Graph. |
| 80 | |
| 81 | Args: |
| 82 | graph: tf.Graph. |
| 83 | run_meta: RunMetadata proto used to complete shape information. |
| 84 | add_trace: Whether to add op trace information. |
| 85 | add_trainable_var: Whether to assign tf.compat.v1.trainable_variables() op |
| 86 | type '_trainable_variables'. |
| 87 | Returns: |
| 88 | logged_ops: dict mapping from op_name to OpLogEntry. |
| 89 | string_to_id: dict mapping from string to id. |
| 90 | """ |
| 91 | if run_meta: |
| 92 | graph = _fill_missing_graph_shape(graph, run_meta) |
| 93 | |
| 94 | op_missing_shape = 0 |
| 95 | logged_ops = {} |
| 96 | string_to_id = {} |
| 97 | string_to_id['none'] = len(string_to_id) |
| 98 | # TODO(xpan): Work with Profiler more efficiently. |
| 99 | for op in graph.get_operations(): |
| 100 | try: |
| 101 | stats = ops.get_stats_for_node_def( |
| 102 | graph, op.node_def, REGISTERED_FLOP_STATS) |
| 103 | except ValueError: |
| 104 | # Catch Exception When shape is incomplete. Skip it. |
| 105 | op_missing_shape += 1 |
| 106 | stats = None |
| 107 | |
| 108 | entry = tfprof_log_pb2.OpLogEntry() |
| 109 | entry.name = op.name |
| 110 | add_entry = False |
| 111 | if stats and stats.value: |
| 112 | entry.float_ops = int(stats.value) |
| 113 | add_entry = True |
| 114 | |
| 115 | if add_trace: |
| 116 | for tb in op.traceback_with_start_lines: |
| 117 | trace = entry.code_def.traces.add() |
| 118 | trace.file_id = _str_id(tb[0], string_to_id) if tb[0] else 0 |
| 119 | trace.lineno = tb[1] if tb[1] else -1 |
| 120 | trace.function_id = _str_id(tb[2], string_to_id) if tb[2] else 0 |
| 121 | trace.line_id = _str_id(tb[3], string_to_id) if tb[3] else 0 |
| 122 | trace.func_start_line = tb[4] if tb[4] else -1 |
| 123 | add_entry = True |
| 124 | |
| 125 | if add_entry: |
| 126 | logged_ops[entry.name] = entry |
| 127 | |
| 128 | if add_trainable_var: |
| 129 | for v in graph.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES): |
| 130 | if v.op.name not in logged_ops: |
| 131 | entry = tfprof_log_pb2.OpLogEntry() |
| 132 | entry.name = v.op.name |
| 133 | entry.types.append(TRAINABLE_VARIABLES) |
| 134 | logged_ops[entry.name] = entry |
no test coverage detected