How to do Relay to Relay transformation by traversing each node?

Development
#1

Hi All,

I’ve written a simple code that traverses Relay program and prints each node, after that I want to construct a NEW RELAY module using the traversed nodes.

Following is the code that I have. However, currently, it is throwing an error. I do not know if the following is the best way to do it, but basically, I’d like to traverse the nodes in Relay program and create another equivalent Relay program using the nodes. If you know a better way of achieving this, please let me know.

# an empty Relay module
mod_new = relay.Module()


def print_and_create_new_module(node):
    """ This function prints node and attempts to create a new relay Module"""
    print("\n")

    # Create a module using the node
    temp_mod = relay.Module.from_expr(node)

    # Update the mod_new with the temp_mod
    mod_new.update(temp_mod)

    print("Current node: {}".format(node))
    print("\n")

def traverse():
    # Create a simple Relay program.
    c = relay.const(2.0, "float32")
    x = relay.var("x", "float32")
    y = relay.add(c, c)

    func = relay.Function([x], y)
    mod = relay.Module({"main": func})

    # traverse the relay IR and create a new RELAY module with every traversed node.
    relay.analysis.post_order_visit(mod['main'], lambda node: print_and_create_new_module(node))

Error:

tvm._ffi.base.TVMError: Traceback (most recent call last): File “C:\anaconda\envs\python3.8.1_test\lib\site-packages\tvm-0.7.dev0-py3.8-win-amd64.egg\tvm_ffi_ctypes\function.py”, line 72, in cfun rv = local_pyfunc(*pyargs) File “C:/repos1/tvm_test/relay_traverse_graph.py”, line 37, in relay.analysis.post_order_visit(mod[‘main’], lambda node: print_and_create_new_module(node)) File “C:/repos1/tvm_test/relay_traverse_graph.py”, line 22, in print_and_create_new_module mod_new.update(temp_mod) File “C:\anaconda\envs\python3.8.1_test\lib\site-packages\tvm-0.7.dev0-py3.8-win-amd64.egg\tvm\relay\module.py”, line 131, in update return _module.Module_Update(self, other) File “C:\anaconda\envs\python3.8.1_test\lib\site-packages\tvm-0.7.dev0-py3.8-win-amd64.egg\tvm_ffi_ctypes\function.py”, line 207, in call raise get_last_ffi_error() File “C:\repos1\tvm\src\ir\module.cc”, line 182 TVMError: Check failed: (*it).second == var (GlobalVar(main) vs. GlobalVar(main)) :

How to get graph representation of Relay program?
#2

I dont think you understand what a module mean. A module hold a bunch of functions. A single expression in a function is not a function, and you should not put it in a module. Beside, there can only be one main function, and when you update(temp_mod), you are adding multiple main function.

1 Like
#3

@MarisaKirisame - you are right. After reading backend code of Module, I understand that what I posted does not make sense.

I believe you mentioned earlier that Relay program is a graph representation. If that is the case I am under the assumption that nodes in relay program are functions?

What do you suggest for getting a graph representation of relay program?

#4

No. Some nodes in the relay program are FunctionNode. They are like lambda expression, but only some of them are. In your example,

def traverse(): # Create a simple Relay program. c = relay.const(2.0, “float32”) x = relay.var(“x”, “float32”) y = relay.add(c, c)

func = relay.Function([x], y)

the y is a computation graph, if that is what you want.

1 Like
#5

Thank you @MarisaKirisame. This is super useful.

So my question is how to traverse and transform the computation graph? Let us say, I’d like to transform the add in my example to multiple? This is what I have been trying to do.

#6

please read expr.py. add is a function call with the operator as the operator add.

#7

Did your questiones all solved?I also need to generate a compute graph from Relay IR which has finished Graph optimization or quantization.