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