# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Graph debug results dumping class.""" import collections import json import os import numpy as np import tvm GRAPH_DUMP_FILE_NAME = '_tvmdbg_graph_dump.json' CHROME_TRACE_FILE_NAME = "_tvmdbg_execution_trace.json" ChromeTraceEvent = collections.namedtuple( 'ChromeTraceEvent', ['ts', 'tid', 'pid', 'name', 'ph'] ) class DebugResult(object): """Graph debug data module. Data dump module manage all the debug data formatting. Output data and input graphs are formatted and dumped to file. Frontend read these data and graph for visualization. Parameters ---------- graph_json : str The graph to be deployed in json format output by graph compiler. Each operator (tvm_op) in the graph will have a one to one mapping with the symbol in libmod which is used to construct a "PackedFunc" . dump_path : str Output data path is read/provided from frontend """ def __init__(self, graph_json, dump_path): self._dump_path = dump_path self._output_tensor_list = [] self._time_list = [] self._parse_graph(graph_json) # dump the json information self.dump_graph_json(graph_json) def _parse_graph(self, graph_json): """Parse and extract the JSON graph and update the nodes, shapes and dltype. Parameters ---------- graph_json : str or graph class The graph to be deployed in json format output by JSON graph. """ json_obj = json.loads(graph_json) self._nodes_list = json_obj['nodes'] self._shapes_list = json_obj['attrs']['shape'] self._dtype_list = json_obj['attrs']['dltype'] self._update_graph_json() def _update_graph_json(self): """update the nodes_list with name, shape and data type, for temporarily storing the output. """ nodes_len = len(self._nodes_list) for i in range(nodes_len): node = self._nodes_list[i] input_list = [] for input_node in node['inputs']: input_list.append(self._nodes_list[input_node[0]]['name']) node['inputs'] = input_list dtype = str("type: " + self._dtype_list[1][i]) if 'attrs' not in node: node['attrs'] = {} node['op'] = "param" else: node['op'] = node['attrs']['func_name'] node['attrs'].update({"T": dtype}) node['shape'] = self._shapes_list[1][i] def _cleanup_tensors(self): """Remove the tensor dump file (graph wont be removed) """ for filename in os.listdir(self._dump_path): if os.path.isfile(filename) and not filename.endswith(".json"): os.remove(filename) def get_graph_nodes(self): """Return the nodes list """ return self._nodes_list def get_graph_node_shapes(self): """Return the nodes shapes list """ return self._shapes_list def get_graph_node_output_num(self, node): """Return the number of outputs of a node """ return 1 if node['op'] == 'param' else int(node['attrs']['num_outputs']) def get_graph_node_dtypes(self): """Return the nodes dtype list """ return self._dtype_list def get_output_tensors(self): """Dump the outputs to a temporary folder, the tensors are in numpy format """ eid = 0 order = 0 output_tensors = {} for node, time in zip(self._nodes_list, self._time_list): num_outputs = self.get_graph_node_output_num(node) for j in range(num_outputs): order += time[0] key = node['name'] + "_" + str(j) output_tensors[key] = self._output_tensor_list[eid] eid += 1 return output_tensors def dump_output_tensor(self): """Dump the outputs to a temporary folder, the tensors are in numpy format """ #cleanup existing tensors before dumping self._cleanup_tensors() eid = 0 order = 0 output_tensors = {} for node, time in zip(self._nodes_list, self._time_list): num_outputs = self.get_graph_node_output_num(node) for j in range(num_outputs): order += time[0] key = node['name'] + "_" + str(j) + "__" + str(order) output_tensors[key] = self._output_tensor_list[eid] eid += 1 with open(os.path.join(self._dump_path, "output_tensors.params"), "wb") as param_f: param_f.write(save_tensors(output_tensors)) def dump_chrome_trace(self): """Dump the trace to the Chrome trace.json format. """ def s_to_us(t): return t * 10 ** 6 starting_times = np.zeros(len(self._time_list) + 1) starting_times[1:] = np.cumsum([times[0] for times in self._time_list]) def node_to_events(node, times, starting_time): return [ ChromeTraceEvent( ts=s_to_us(starting_time), tid=1, pid=1, ph='B', name=node['name'], ), ChromeTraceEvent( # Use start + duration instead of end to ensure precise timings. ts=s_to_us(times[0] + starting_time), tid=1, pid=1, ph='E', name=node['name'], ), ] events = [ e for (node, times, starting_time) in zip( self._nodes_list, self._time_list, starting_times) for e in node_to_events(node, times, starting_time)] result = dict( displayTimeUnit='ns', traceEvents=[e._asdict() for e in events] ) with open(os.path.join(self._dump_path, CHROME_TRACE_FILE_NAME), "w") as trace_f: json.dump(result, trace_f) def dump_graph_json(self, graph): """Dump json formatted graph. Parameters ---------- graph : json format json formatted JSON graph contain list of each node's name, shape and type. """ graph_dump_file_name = GRAPH_DUMP_FILE_NAME with open(os.path.join(self._dump_path, graph_dump_file_name), 'w') as outfile: json.dump(graph, outfile, indent=4, sort_keys=False) def get_debug_result(self, sort_by_time=True): """Return the debugger result""" header = ["Node Name", "Ops", "Time(us)", "Time(%)", "Shape", "Inputs", "Outputs"] lines = ["---------", "---", "--------", "-------", "-----", "------", "-------"] eid = 0 data = [] total_time = sum(time[0] for time in self._time_list) for node, time in zip(self._nodes_list, self._time_list): num_outputs = self.get_graph_node_output_num(node) for j in range(num_outputs): op = node['op'] if node['op'] == 'param': eid += 1 continue name = node['name'] shape = str(self._output_tensor_list[eid].shape) time_us = round(time[0] * 1000000, 3) time_percent = round(((time[0] / total_time) * 100), 3) inputs = str(node['attrs']['num_inputs']) outputs = str(node['attrs']['num_outputs']) node_data = [name, op, time_us, time_percent, shape, inputs, outputs] data.append(node_data) eid += 1 if sort_by_time: # Sort on the basis of execution time. Prints the most expensive ops in the start. data = sorted(data, key=lambda x: x[2], reverse=True) # Insert a row for total time at the end. rounded_total_time = round(total_time * 1000000, 3) data.append(["Total_time", "-", rounded_total_time, "-", "-", "-", "-", "-"]) fmt = "" for i, _ in enumerate(header): max_len = len(header[i]) for j, _ in enumerate(data): item_len = len(str(data[j][i])) if item_len > max_len: max_len = item_len fmt = fmt + "{:<" + str(max_len + 2) + "}" log = [fmt.format(*header)] log.append(fmt.format(*lines)) for row in data: log.append(fmt.format(*row)) return '\n'.join(log) def display_debug_result(self, sort_by_time=True): """Displays the debugger result""" print(self.get_debug_result(sort_by_time)) def save_tensors(params): """Save parameter dictionary to binary bytes. The result binary bytes can be loaded by the GraphModule with API "load_params". Parameters ---------- params : dict of str to NDArray The parameter dictionary. Returns ------- param_bytes: bytearray Serialized parameters. """ _save_tensors = tvm.get_global_func("tvm.relay._save_param_dict") args = [] for k, v in params.items(): args.append(k) args.append(tvm.nd.array(v)) return _save_tensors(*args)