def test_save_debug_info_enabled(self): root = tracking.AutoTrackable() root.f = def_function.function( lambda x: math_ops.mul(2., x, name="DEBUG_INFO_OP"), input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) save_dir = os.path.join(self.get_temp_dir(), "saved_model") save.save( root, save_dir, root.f, options=save_options.SaveOptions(save_debug_info=True)) debug_info_file_name = os.path.join(save_dir, "debug", "saved_model_debug_info.pb") self.assertTrue(os.path.exists(debug_info_file_name)) debug_info = graph_debug_info_pb2.GraphDebugInfo() with open(debug_info_file_name, "rb") as f: debug_info.ParseFromString(f.read()) # Verify that there is a trace for DEBUG_INFO_OP just to ensure that # function debug info tracing is nominally functioning. found_op = False for key in debug_info.traces.keys(): if key.startswith("DEBUG_INFO_OP@"): found_op = True break self.assertTrue(found_op, "Did not find DEBUG_INFO_OP in trace")
def parse_saved_model_with_debug_info(export_dir): """Reads the savedmodel as well as the graph debug info. Args: export_dir: Directory containing the SavedModel and GraphDebugInfo files. Returns: `SavedModel` and `GraphDebugInfo` protocol buffers. Raises: IOError: If the saved model file does not exist, or cannot be successfully parsed. Missing graph debug info file is fine. """ saved_model = _parse_saved_model(export_dir) debug_info_path = os.path.join(saved_model_utils.get_debug_dir(export_dir), constants.DEBUG_INFO_FILENAME_PB) debug_info = graph_debug_info_pb2.GraphDebugInfo() if file_io.file_exists(debug_info_path): with file_io.FileIO(debug_info_path, "rb") as debug_file: try: debug_info.ParseFromString(debug_file.read()) except message.DecodeError as e: raise IOError(f"Cannot parse file {debug_info_path}: {e}.") return (saved_model, debug_info)
def create_graph_debug_info_def(func_named_operations): """Construct and returns a `GraphDebugInfo` protocol buffer. Args: func_named_operations: An iterable of (func_name, op.Operation) tuples where the Operation instances have a _traceback members. The func_name should be the empty string for operations in the top-level Graph. Returns: GraphDebugInfo protocol buffer. Raises: TypeError: If the arguments are not of the correct proto buffer type. """ # Creates an empty GraphDebugInfoDef proto. graph_debug_info_def = graph_debug_info_pb2.GraphDebugInfo() # Gets the file names and line numbers for the exported node names. Also # collects the unique file names. all_file_names = set() node_to_trace = {} for func_name, op in func_named_operations: try: op_traceback = op.traceback except AttributeError: # Some ops synthesized on as part of function or control flow definition # do not have tracebacks. continue # Gets the stack trace of the operation and then the file location. node_name = op.name + "@" + func_name node_to_trace[node_name] = _compute_useful_frames(op_traceback, 10) for frame in node_to_trace[node_name]: all_file_names.add(frame.filename) # Sets the `files` field in the GraphDebugInfo proto graph_debug_info_def.files.extend(all_file_names) # Builds a mapping between file names and index of the `files` field, so we # only store the indexes for the nodes in the GraphDebugInfo. file_to_index = dict( [(y, x) for x, y in enumerate(graph_debug_info_def.files)]) # Creates the FileLineCol proto for each node and sets the value in the # GraphDebugInfo proto. We only store the file name index for each node to # save the storage space. for node_name, frames in node_to_trace.items(): trace_def = graph_debug_info_def.traces[node_name] for frame in reversed(frames): trace_def.file_line_cols.add( file_index=file_to_index[frame.filename], line=frame.lineno) return graph_debug_info_def
def adjust_debug_info_func_names(self, debug_info): """Rewrite func names in the debug info by using the concrete func names.""" output_debug_info = graph_debug_info_pb2.GraphDebugInfo() output_debug_info.files[:] = debug_info.files for key in debug_info.traces: node, func = key.split("@") new_func = "" if func in self._concrete_functions: new_func = self._concrete_functions[func].function_def.signature.name output_debug_info.traces[node + "@" + new_func].CopyFrom( debug_info.traces[key]) return output_debug_info
def f(original_nodes): """Function to create `GraphDebugInfo` for the given `original_nodes`.""" if not saved_debug_info: return None output_debug_info = graph_debug_info_pb2.GraphDebugInfo() # All the files are copied over, so the index wouldn't be changed. output_debug_info.files[:] = saved_debug_info.files # We only copy over the debug info for the input nodes for func, node in original_nodes: debug_key = node + "@" + func output_debug_info.traces[debug_key].CopyFrom( saved_debug_info.traces[debug_key]) return output_debug_info
def create_graph_debug_info_def(operations): """Construct and returns a `GraphDebugInfo` protocol buffer. Args: operations: An iterable of op.Operation objects having _traceback members. Returns: GraphDebugInfo protocol buffer. Raises: TypeError: If the arguments are not of the correct proto buffer type. """ # Creates an empty GraphDebugInfoDef proto. graph_debug_info_def = graph_debug_info_pb2.GraphDebugInfo() # Gets the file names and line numbers for the exported node names. Also # collects the unique file names. all_file_names = set() node_to_trace = {} for op in operations: # Gets the stack trace of the operation and then the file location. node_name = op.name node_to_trace[node_name] = error_interpolation.compute_useful_frames( op, 10) for frame in node_to_trace[node_name]: all_file_names.add(frame[tf_stack.TB_FILENAME]) # Sets the `files` field in the GraphDebugInfo proto graph_debug_info_def.files.extend(all_file_names) # Builds a mapping between file names and index of the `files` field, so we # only store the indexes for the nodes in the GraphDebugInfo. file_to_index = dict([(y, x) for x, y in enumerate(graph_debug_info_def.files)]) # Creates the FileLineCol proto for each node and sets the value in the # GraphDebugInfo proto. We only store the file name index for each node to # save the storage space. for node_name, frames in node_to_trace.items(): trace_def = graph_debug_info_def.traces[node_name] for frame in reversed(frames): trace_def.file_line_cols.add( file_index=file_to_index[frame[tf_stack.TB_FILENAME]], line=frame[tf_stack.TB_LINENO], func=frame[tf_stack.TB_FUNCNAME], code=frame[tf_stack.TB_CODEDICT]) return graph_debug_info_def