Retrieve AST of a network


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]




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

The lower function is called in


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


DoLower calls into GetScheduleArgs, which will call TOPI schedule of opperators (the schedules are registered in nnvm/python/nnvm/top)