Iterate through layers in the same order as the forward pass

Is there a mechanism to iterate through the modules of a network in the exact order they’re executed in the forward pass? I’m looking for a solution that does not require the usage of torch.nn.Sequential.

Because forward can be customized by users, there is no general way guaranteed to work on arbitrary networks.

One possibility is to run the forward once, and use jit tracer to trace the graph. But you may get lower-level ops that implement each layer rather than the layers as individual ops (e.g., instance norm may become: reshape + batch_norm + reshape + mul + add)

Thanks for your reply, SimonW. I ended up replacing layers in the network itself.

1 Like