1# Licensed to the Apache Software Foundation (ASF) under one 2# or more contributor license agreements. See the NOTICE file 3# distributed with this work for additional information 4# regarding copyright ownership. The ASF licenses this file 5# to you under the Apache License, Version 2.0 (the 6# "License"); you may not use this file except in compliance 7# with the License. You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, 12# software distributed under the License is distributed on an 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14# KIND, either express or implied. See the License for the 15# specific language governing permissions and limitations 16# under the License. 17"""Graph debug results dumping class.""" 18import collections 19import json 20import os 21import numpy as np 22import tvm 23 24GRAPH_DUMP_FILE_NAME = '_tvmdbg_graph_dump.json' 25CHROME_TRACE_FILE_NAME = "_tvmdbg_execution_trace.json" 26 27ChromeTraceEvent = collections.namedtuple( 28 'ChromeTraceEvent', 29 ['ts', 'tid', 'pid', 'name', 'ph'] 30) 31 32 33class DebugResult(object): 34 """Graph debug data module. 35 36 Data dump module manage all the debug data formatting. 37 Output data and input graphs are formatted and dumped to file. 38 Frontend read these data and graph for visualization. 39 40 Parameters 41 ---------- 42 graph_json : str 43 The graph to be deployed in json format output by nnvm graph. Each operator (tvm_op) 44 in the graph will have a one to one mapping with the symbol in libmod which is used 45 to construct a "PackedFunc" . 46 47 dump_path : str 48 Output data path is read/provided from frontend 49 """ 50 51 def __init__(self, graph_json, dump_path): 52 self._dump_path = dump_path 53 self._output_tensor_list = [] 54 self._time_list = [] 55 self._parse_graph(graph_json) 56 # dump the json information 57 self.dump_graph_json(graph_json) 58 59 def _parse_graph(self, graph_json): 60 """Parse and extract the NNVM graph and update the nodes, shapes and dltype. 61 62 Parameters 63 ---------- 64 graph_json : str or graph class 65 The graph to be deployed in json format output by nnvm graph. 66 """ 67 json_obj = json.loads(graph_json) 68 self._nodes_list = json_obj['nodes'] 69 self._shapes_list = json_obj['attrs']['shape'] 70 self._dtype_list = json_obj['attrs']['dltype'] 71 self._update_graph_json() 72 73 def _update_graph_json(self): 74 """update the nodes_list with name, shape and data type, 75 for temporarily storing the output. 76 """ 77 78 nodes_len = len(self._nodes_list) 79 for i in range(nodes_len): 80 node = self._nodes_list[i] 81 input_list = [] 82 for input_node in node['inputs']: 83 input_list.append(self._nodes_list[input_node[0]]['name']) 84 node['inputs'] = input_list 85 dtype = str("type: " + self._dtype_list[1][i]) 86 if 'attrs' not in node: 87 node['attrs'] = {} 88 node['op'] = "param" 89 else: 90 node['op'] = node['attrs']['func_name'] 91 node['attrs'].update({"T": dtype}) 92 node['shape'] = self._shapes_list[1][i] 93 94 def _cleanup_tensors(self): 95 """Remove the tensor dump file (graph wont be removed) 96 """ 97 for filename in os.listdir(self._dump_path): 98 if os.path.isfile(filename) and not filename.endswith(".json"): 99 os.remove(filename) 100 101 def get_graph_nodes(self): 102 """Return the nodes list 103 """ 104 return self._nodes_list 105 106 def get_graph_node_shapes(self): 107 """Return the nodes shapes list 108 """ 109 return self._shapes_list 110 111 def get_graph_node_output_num(self, node): 112 """Return the number of outputs of a node 113 """ 114 return 1 if node['op'] == 'param' else int(node['attrs']['num_outputs']) 115 116 def get_graph_node_dtypes(self): 117 """Return the nodes dtype list 118 """ 119 return self._dtype_list 120 121 def get_output_tensors(self): 122 """Dump the outputs to a temporary folder, the tensors are in numpy format 123 """ 124 eid = 0 125 order = 0 126 output_tensors = {} 127 for node, time in zip(self._nodes_list, self._time_list): 128 num_outputs = self.get_graph_node_output_num(node) 129 for j in range(num_outputs): 130 order += time[0] 131 key = node['name'] + "_" + str(j) 132 output_tensors[key] = self._output_tensor_list[eid] 133 eid += 1 134 return output_tensors 135 136 def dump_output_tensor(self): 137 """Dump the outputs to a temporary folder, the tensors are in numpy format 138 """ 139 #cleanup existing tensors before dumping 140 self._cleanup_tensors() 141 eid = 0 142 order = 0 143 output_tensors = {} 144 for node, time in zip(self._nodes_list, self._time_list): 145 num_outputs = self.get_graph_node_output_num(node) 146 for j in range(num_outputs): 147 order += time[0] 148 key = node['name'] + "_" + str(j) + "__" + str(order) 149 output_tensors[key] = self._output_tensor_list[eid] 150 eid += 1 151 152 with open(os.path.join(self._dump_path, "output_tensors.params"), "wb") as param_f: 153 param_f.write(save_tensors(output_tensors)) 154 155 def dump_chrome_trace(self): 156 """Dump the trace to the Chrome trace.json format. 157 """ 158 def s_to_us(t): 159 return t * 10 ** 6 160 161 starting_times = np.zeros(len(self._time_list) + 1) 162 starting_times[1:] = np.cumsum([times[0] for times in self._time_list]) 163 164 def node_to_events(node, times, starting_time): 165 return [ 166 ChromeTraceEvent( 167 ts=s_to_us(starting_time), 168 tid=1, 169 pid=1, 170 ph='B', 171 name=node['name'], 172 ), 173 ChromeTraceEvent( 174 # Use start + duration instead of end to ensure precise timings. 175 ts=s_to_us(times[0] + starting_time), 176 tid=1, 177 pid=1, 178 ph='E', 179 name=node['name'], 180 ), 181 ] 182 events = [ 183 e for (node, times, starting_time) in zip( 184 self._nodes_list, self._time_list, starting_times) 185 for e in node_to_events(node, times, starting_time)] 186 result = dict( 187 displayTimeUnit='ns', 188 traceEvents=[e._asdict() for e in events] 189 ) 190 191 with open(os.path.join(self._dump_path, CHROME_TRACE_FILE_NAME), "w") as trace_f: 192 json.dump(result, trace_f) 193 194 def dump_graph_json(self, graph): 195 """Dump json formatted graph. 196 197 Parameters 198 ---------- 199 graph : json format 200 json formatted NNVM graph contain list of each node's 201 name, shape and type. 202 """ 203 graph_dump_file_name = GRAPH_DUMP_FILE_NAME 204 with open(os.path.join(self._dump_path, graph_dump_file_name), 'w') as outfile: 205 json.dump(graph, outfile, indent=4, sort_keys=False) 206 207 def display_debug_result(self, sort_by_time=True): 208 """Displays the debugger result" 209 """ 210 header = ["Node Name", "Ops", "Time(us)", "Time(%)", "Shape", "Inputs", "Outputs"] 211 lines = ["---------", "---", "--------", "-------", "-----", "------", "-------"] 212 eid = 0 213 data = [] 214 total_time = sum(time[0] for time in self._time_list) 215 for node, time in zip(self._nodes_list, self._time_list): 216 num_outputs = self.get_graph_node_output_num(node) 217 for j in range(num_outputs): 218 op = node['op'] 219 if node['op'] == 'param': 220 eid += 1 221 continue 222 name = node['name'] 223 shape = str(self._output_tensor_list[eid].shape) 224 time_us = round(time[0] * 1000000, 3) 225 time_percent = round(((time[0] / total_time) * 100), 3) 226 inputs = str(node['attrs']['num_inputs']) 227 outputs = str(node['attrs']['num_outputs']) 228 node_data = [name, op, time_us, time_percent, shape, inputs, outputs] 229 data.append(node_data) 230 eid += 1 231 232 if sort_by_time: 233 # Sort on the basis of execution time. Prints the most expensive ops in the start. 234 data = sorted(data, key=lambda x: x[2], reverse=True) 235 # Insert a row for total time at the end. 236 rounded_total_time = round(total_time * 1000000, 3) 237 data.append(["Total_time", "-", rounded_total_time, "-", "-", "-", "-", "-"]) 238 239 fmt = "" 240 for i, _ in enumerate(header): 241 max_len = len(header[i]) 242 for j, _ in enumerate(data): 243 item_len = len(str(data[j][i])) 244 if item_len > max_len: 245 max_len = item_len 246 fmt = fmt + "{:<" + str(max_len + 2) + "}" 247 print(fmt.format(*header)) 248 print(fmt.format(*lines)) 249 for row in data: 250 print(fmt.format(*row)) 251 252def save_tensors(params): 253 """Save parameter dictionary to binary bytes. 254 255 The result binary bytes can be loaded by the 256 GraphModule with API "load_params". 257 258 Parameters 259 ---------- 260 params : dict of str to NDArray 261 The parameter dictionary. 262 263 Returns 264 ------- 265 param_bytes: bytearray 266 Serialized parameters. 267 """ 268 _save_tensors = tvm.get_global_func("_save_param_dict") 269 270 args = [] 271 for k, v in params.items(): 272 args.append(k) 273 args.append(tvm.nd.array(v)) 274 return _save_tensors(*args) 275