[WIP] Relay visualization: exporter + visualizer #4370
Conversation
| @@ -22,11 +22,14 @@ | |||
| """ | |||
| from . import _analysis | |||
| from . import _make | |||
| from .expr import Expr | |||
| from .expr import Expr, Function, Var, Call, TupleGetItem | |||
| 'ret_type': { | ||
| 'dtype': node.ret_type.dtype, | ||
| 'shape': [int(x) for x in node.ret_type.shape] | ||
| } |
There was a problem hiding this comment.
How about there is a type attribute on every op and it contains a string that describes the type of what the node yields? It can be an empty string if there is no value. Any visualizer that wants to actually understand the structure of e.g. types can just deal with Relay directly.
There was a problem hiding this comment.
@broune I don't see a way to obtain the return type of the op. In fact, it looks like the "ret_type" field in the Function node is optional. It appears that type inference occurs when the Relay program gets compiled. Let me go ahead and remove this field.
There was a problem hiding this comment.
One of the main use cases for this is to print out the state of the IR as it gets mutated through compilation, where type inference is one of the first things (IIRC actually the first thing) that happens, so it is good to show the type if it is present - that accurately reflects the IR.
There was a problem hiding this comment.
I'll have to ask TVM experts for help. Specifically: 1) where type inference occurs, and 2) how to get the return type of each Op.
| relayviz_nodes.append({ | ||
| 'node_kind': 'Call', | ||
| 'op': node_dict[node.op], | ||
| 'args': [node_dict[arg] for arg in node.args] |
There was a problem hiding this comment.
Can we call the within-graph operands of all nodes the same regardless of what kind of node it is? E.g. operands.
| 'node_kind': 'Call', | ||
| 'op': node_dict[node.op], | ||
| 'args': [node_dict[arg] for arg in node.args] | ||
| }) |
There was a problem hiding this comment.
Wouldn't you want this to also have a type field that says what type is yielded by the call?
There was a problem hiding this comment.
See my comment above about inability to obtain the return type of an op.
| elif isinstance(node, TupleGetItem): | ||
| relayviz_nodes.append({ | ||
| 'node_kind': 'TupleGetItem', | ||
| 'tuple_value': node_dict[node.tuple_value], |
There was a problem hiding this comment.
This is also jus the operand of this operation, so could be called with a consistent name to other ops.
| relayviz_nodes.append({ | ||
| 'node_kind': 'TupleGetItem', | ||
| 'tuple_value': node_dict[node.tuple_value], | ||
| 'index': node.index |
There was a problem hiding this comment.
Is this the same as an attribute on an op? Why does this node type need to be special? (actually, why does any node type need to be special?)
There was a problem hiding this comment.
@broune Relay is not a dataflow graph but a programming language. So the current prototype makes differentiation between different kinds of language constructs.
@junrushao1994 What's your take on this? Should we collapse Relay nodes into generic universal node type?
There was a problem hiding this comment.
This is about providing sufficient information to the visualizer so that it can create a good visualization while being as simple as possible. In this case, I wonder whether the visualizer has a good use of a special case for the index field or TupleGetItem versus other ops. If not, seems simpler to not have the special case for this, and same question on a case-by-case basis for other cases.
There was a problem hiding this comment.
@broune I see your point. In this case, you're okay with grouping Relay nodes into, say, 3-4 possible categories?
There was a problem hiding this comment.
Sorry for the late reply. I'm OK with whatever you think makes sense on this; you're more involved in the details of this than I am, and what I most want (and am very OK with!) is to have the ability to see Relay graphs easily visualized. I just wanted to raise this as a thing to consider.
| relayviz_nodes.append({ | ||
| 'node_kind': 'Op', | ||
| 'name': node.name, | ||
| 'attrs': {} |
There was a problem hiding this comment.
Note to myself: it is actually possible to obtain attributes of an op using Python.
|
close due to inactive status, @hcho3 feel free to reopen :) |
Implements the RFC https://discuss.tvm.ai/t/rfc-visualizing-relay-program-as-graph/4825/9. Note that this implementation is not yet complete; I put it up early so that the community may start getting the feel of how the visualization feature will be organized.