I would like to dynamically replace the forward function of some layers in a neural network with some other forward functions.
Currently, I achieve this by defining forward hooks on the target layers such that the hook function calls the desired forward, which replaces the original output. However, this implementation has two shortcomings. First, both the original and new forward functions are called, which results in slow-down and waste of resources. Second, using hooks seems to slow down the code significantly.
I was thinking that an alternative solution would be to find all the forward functions that are called, replace some of them, and call the new high-level forward function. However, I am not sure how to find the order in which all forward functions are called.
Can you please let me know if there is an easy way of finding all the forward functions that are called with their correct order?