Retrieve variables and expressions from a Relay function

Questions
1 
data = relay.var("data", relay.TensorType((relay.Any(), 3, 128, 128), "float32"))
conv = relay.nn.conv2d(
    data=data, weight=relay.var("weight"), kernel_size=(3, 3), channels=3, padding=(1, 1)
)
net = relay.Function(relay.analysis.free_vars(conv), conv))

By doing so, we created a Relay function.

I wonder:

  1. How can we retrieve data and conv variable/expression given net?
  2. Given conv, is it possible to retrieve a constructor like lambda x : relay.nn.conv2d(data=x, weight=relay.var("weight"), ...) (i.e., change the input variable data to something else in conv expression).
2

Hi @ganler, for your first question, you can get data and conv by calling net.params[0] and net.body. You can check out FunctionNode implementation in Relay AST here: tvm/function.h at main · apache/tvm · GitHub.

3

@yuchenj Thanks for the reply. Yes, I understand that relay.Function has 3 useful APIs: params, body, ret_type and that .body() is actually the return node/variable of the current function and we need to perform recursive Visit on it to perform either mutation or some kind of analysis.

I was wondering if there’s some graph-level APIs like ONNX (i.e., protobuf) so that we can access the variables/expression by a name or key. I know that this can be achieved by the visitor pattern (but I guess I have to dive into C++ to do so), but I am just curious if there’s also well-supported Python APIs. :blush:

4

@ganler, I think there are no such python APIs to do the access and mutation today, and I think it is good to have. :slight_smile:

1 Like