How to see different IR(Relay/TE/TIR) from my own pytorch model?

Questions
#1

Hello tvm community.

I am learning tvm from the tutorials and other online blogs. I learned that the compile flow in tvm is like the picture(map) below.

Untitled
Untitled995×603 31.3 KB

Then I want to use my own network model to do the whole flow one step by one step. But I noticed that most of the tutorials start from the Relay IR written by the writter, not from the network model. So I wonder how to compile one step by one step like the map shows with my onw model.

Here is what I did and what I want to do. First I built my easy pytorch model with only 1 conv+1 relu+1 MaxPool. Then I changed it into onnx form in case I can see the graph in Netron. After that, I use the code below by following the tutorials.

mod, params = relay.frontend.from_onnx(my_onnx_model, shape_dict,)

print(mod)

I thought the ‘mod’ I printed is the Relay IR. Then I noticed that in the tutorials, it just use

tvm.transform.PassContext)(opt_level=3):

lib = relay.build(mod, target=target,params=params)

dev = tvm.device(str(target,0)) # the target is ‘llvm’

module=tvm.contrib.graph_executor.GraphModule(lib"default")

Therefore I can’t see the other IR forms of my network model. I wonder how I can see them, and whether I can optimize the model like using fusion or other ways. Thank you!

#2

I think you wanna to see the intermediate during build stage and what the pass did. Is that ? please refer to Output all pass names and ir

#3

To see the optimized relayIR, you can use sth like this:

with tvm.transform.PassContext(opt_level=3):
    opt_mod = relay.optimize(mod, target=target, params=params)

To see the lowered final IR form, maybe you can use tvm.lower function, refering tvm.driver — tvm 0.11.dev0 documentation, however I have only used it on single ops, not sure whether it can be applied to a complete model

# declare a matrix element-wise multiply
A = te.placeholder((m, n), name="A")
B = te.placeholder((m, n), name="B")
C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name="C")

s = te.create_schedule([C.op])
# lower will transform the computation from definition to the real
# callable function. With argument `simple_mode=True`, it will
# return you a readable C like statement, we use it here to print the
# schedule result.
print(tvm.lower(s, [A, B, C], simple_mode=True))
#4

Thanks for your answer. I got what I want according to your advice!

#5

Thanks you mgeek. I saw this in the tutorial, too. Actually I want to see the process how tvm compile my model auto, and I have no idea how to use tvm.driver.lower with my net model (lol).Your answer is helpful by suggesting me using tvm.driver for I can use it in another part of my work.

1 Like