Visualizes an independent graph, or a graph executor.
(graph, name_prefix='', pb_graph=None, executors_it=None)
| 26 | |
| 27 | |
| 28 | def visualize(graph, name_prefix='', pb_graph=None, executors_it=None): |
| 29 | """Visualizes an independent graph, or a graph executor.""" |
| 30 | value_map = {} |
| 31 | pb_graph = pb_graph or graph_pb2.GraphDef() |
| 32 | |
| 33 | if isinstance(graph, torch._C.GraphExecutorState): |
| 34 | visualize_graph_executor(graph, name_prefix, pb_graph, |
| 35 | partial(visualize, pb_graph=pb_graph)) |
| 36 | return pb_graph |
| 37 | |
| 38 | # Set up an input node |
| 39 | input_node = pb_graph.node.add(op='input', name=name_prefix + 'input') |
| 40 | for i, value in enumerate(graph.param_node().outputs()): |
| 41 | value_map[value.unique()] = name_prefix + 'input:' + str(i) |
| 42 | |
| 43 | visualize_rec(graph, value_map, name_prefix, pb_graph, executors_it) |
| 44 | |
| 45 | # Gather all outputs |
| 46 | return_node = pb_graph.node.add(op='output', name=name_prefix + 'output') |
| 47 | for value in graph.return_node().inputs(): |
| 48 | return_node.input.append(value_map[value.unique()]) |
| 49 | |
| 50 | return pb_graph |
| 51 | |
| 52 | |
| 53 | def visualize_graph_executor(state, name_prefix, pb_graph, inline_graph): |
no test coverage detected
searching dependent graphs…