MCPcopy Create free account
hub / github.com/pytorch/pytorch / GetGraphInJson

Function GetGraphInJson

caffe2/python/net_drawer.py:264–319  ·  view source on GitHub ↗
(operators_or_net, output_filepath)

Source from the content-addressed store, hash-verified

262
263
264def GetGraphInJson(operators_or_net, output_filepath):
265 operators, _ = _rectify_operator_and_name(operators_or_net, None)
266 blob_strid_to_node_id = {}
267 node_name_counts = defaultdict(int)
268 nodes = []
269 edges = []
270 for op_id, op in enumerate(operators):
271 op_label = op.name + '/' + op.type if op.name else op.type
272 op_node_id = len(nodes)
273 nodes.append({
274 'id': op_node_id,
275 'label': op_label,
276 'op_id': op_id,
277 'type': 'op'
278 })
279 for input_name in op.input:
280 strid = _escape_label(
281 input_name + str(node_name_counts[input_name]))
282 if strid not in blob_strid_to_node_id:
283 input_node = {
284 'id': len(nodes),
285 'label': input_name,
286 'type': 'blob'
287 }
288 blob_strid_to_node_id[strid] = len(nodes)
289 nodes.append(input_node)
290 else:
291 input_node = nodes[blob_strid_to_node_id[strid]]
292 edges.append({
293 'source': blob_strid_to_node_id[strid],
294 'target': op_node_id
295 })
296 for output_name in op.output:
297 strid = _escape_label(
298 output_name + str(node_name_counts[output_name]))
299 if strid in blob_strid_to_node_id:
300 # we are overwriting an existing blob. need to update the count.
301 node_name_counts[output_name] += 1
302 strid = _escape_label(
303 output_name + str(node_name_counts[output_name]))
304
305 if strid not in blob_strid_to_node_id:
306 output_node = {
307 'id': len(nodes),
308 'label': output_name,
309 'type': 'blob'
310 }
311 blob_strid_to_node_id[strid] = len(nodes)
312 nodes.append(output_node)
313 edges.append({
314 'source': op_node_id,
315 'target': blob_strid_to_node_id[strid]
316 })
317
318 with open(output_filepath, 'w') as f:
319 json.dump({'nodes': nodes, 'edges': edges}, f)
320
321

Callers

nothing calls this directly

Calls 4

_escape_labelFunction · 0.85
appendMethod · 0.45
dumpMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…