How to find the order of calls to different forward methods when calling model(input)?

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

ptrblck: <p>How large is the slow down using hooks?<br>
Would creating custom modules by deriving from the corresponding base class and reimplementing the <code>forward</code> work?<br>
You could monkey-patch the <code>forward</code> method, if you want to change it for all instances of the module, but that’s of course not the cleanest way.</p>

MahdiNazemi: <p>I don’t have the exact numbers yet. But on a simple MLP trained on MNIST, setting the batch size to 10,000 leads to a runtime of 25 seconds while setting the batch size to 256 leads to a runtime of 9 minutes. I guess this is mostly due to the overhead of <code>hook</code>s.</p>
<p>That might work. But sometimes I need to replace multiple consecutive layers with a single layer. Currently, in those cases, the forward hook of the first layer in the group stores its inputs, the new layer takes those inputs and processes them, and the forward hook of the last layer in the group modifies the original output to use the new layer’s output.</p>
<p>Ideally, I would like to eliminate all the calls to the old layers to save computation time and achieve speed up.</p>

ptrblck: <p>The timing sounds weird.<br>
Are you sure you are seeing this slowdown from hooks?<br>
I would assume to see some overhead, but yours seems to be really large.</p>
<p>I’m seeing a ~4% slowdown using this dummy code:</p>
<pre><code class="lang-python">model = models.resnet50()
model.cuda()

x = torch.randn(32, 3, 224, 224).cuda()

# warmup
for _ in range(50):
    out = model(x)

nb_iters = 100    
torch.cuda.synchronize()
t0 = time.time()

for _ in range(nb_iters):
    out = model(x)
    
torch.cuda.synchronize()
t1 = time.time()

print('{}s/iter'.format((t1 - t0)/nb_iters))


# Use hooks
def my_hook(m, input, output):
    output = output * 100
    return output
    
for child in model.children():
    child.register_forward_hook(my_hook)


# warmup
for _ in range(50):
    out = model(x)

nb_iters = 100    
torch.cuda.synchronize()
t0 = time.time()

for _ in range(nb_iters):
    out = model(x)
    
torch.cuda.synchronize()
t1 = time.time()

print('{}s/iter'.format((t1 - t0)/nb_iters))
</code></pre>

MahdiNazemi: <p>Thank you for sharing this snippet of code. It is really helpful.</p>
<p>I profiled the code and it seems that the issue is not due to the <code>hook</code>s, but maybe the function I am replacing. Regardless of whether the batch size is 256 or 10,000, the first layer that I replace takes 12 seconds on GPU and the second one takes about 6 seconds. As a result, the configuration that has a large number of small batches takes a lot longer.</p>
<p>Most of my functions include a single line, which is very long and includes bit-wise operations on <code>bool</code> tensors, e.g.<br>
<code>out[:, 0] = (x[:, 0] &amp; ~x[:, 1]) | (x[:, 124] &amp; (x[:, 118] &amp; ~x[:, 254]))...</code></p>
<p>In both cases, the GPU utilization is around 20%.</p>