Retrieve AST of a network


#1

How do I obtain the IR of a network written in NNVM for custom IR optimization?
That is, assume I have the following network

net, params = nnvm.testing.resnet.get_workload(num_layers=18, batch_size=batch_size, image_shape=image_shape)

and I would like to obtain the IR for writing custom optimizations as described in the tutorial here [1]

[1] https://docs.tvm.ai/tutorials/dev/low_level_custom_pass.html

Thanks,


#2

You can hook nnvm.compiler.lower by calling tvm.register_func with override=True

The lower function is called in


#3

Thanks, but I’m still not clear how to generate a schedule of the network to call the ‘lower’ function.


#4

DoLower calls into GetScheduleArgs https://github.com/dmlc/tvm/blob/bbc78221777a06e3d6b6ee0c07f51f9f37201928/nnvm/src/compiler/compile_engine.cc#L141, which will call TOPI schedule of opperators (the schedules are registered in nnvm/python/nnvm/top)