Leverage torchinfo to easily use targeted hooks?

[eliegoudout]

Dear TylerYep, all,

First of all, thanks for your work, itlooks amazing and is very easy to use :)
I have a question regarding hooks. Imagine I have the following silly network.

class Silly(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax()

    def forward(self, x):
        x = self.relu(self.fc(x))
        x = self.relu(self.fc(x))
        x = self.softmax(self.fc(x))
        return x

It is summarized as follows.

>>> torchinfo.summary(Silly(), (10,), batch_dim=0)
===================================================
Layer (type:depth-idx)  Output Shape  Param #
===================================================
Silly                   [1, 10]       --
├─Linear: 1-1           [1, 10]       110
├─ReLU: 1-2             [1, 10]       --
├─Linear: 1-3           [1, 10]       (recursive)
├─ReLU: 1-4             [1, 10]       --
├─Linear: 1-5           [1, 10]       (recursive)
├─Softmax: 1-6          [1, 10]       --
===================================================

Now, imagine I would like to hook forward the output of Linear: 1-5. Is there a simple way to do so using your package?

The reason why I'm thinking about leveraging your package because it seems that directly using torch.register_forward_hook won't work. Specifically, if I hook Silly().fc directly, then it will execute 3 times for every call of Silly(), which is not wanted. On the other hand, it feels like your package does most of the work by understanding the true architecture of a forward call and differentiating between Linear: 1-1, Linear: 1-3 and Linear: 1-5. This seems useful if I want to target a specific layer based on the summary output.

Thanks in advance!

Cheers.

[eliegoudout]

Also, this would mean being able to setup hooks at runtime, out of the net definition, which looks very attractive to me.

[eliegoudout]

I think I got a solution, but it looks weak / not very robust, and I would appreciate any feedback.
For a given net, I force the user to provide an input_shape so I can save

summ = torchinfo.summary(net, input_shape)

Than, I turn to summ.summary_list, which is a list of LayerInfo instances, representing every layer found by torchinfo on a forward pass.
In the Silly example I provided for example, I would have (written by hand on a computer without python, so it might be an approximative reconstitution)

>>> summ.summary_list
[Silly: 0, Linear: 1, ReLU: 1, Linear: 1, ReLU: 1, Linear: 1, Softmax: 1]

To be clear (correct me if I'm wrong), every LayerInfo instance points to the corresponding nn.Module and different LayerInfo might (and do in this example) point to the same nn.Module.

From this, imagine I would like to save the activations of Linear: 1-5. Then, I leverage the fact that the corresponding nn.Module is called 3 times for every net call and that the third time correspond to Linear: 1-5.

But is this approach robust?? Can I always assume, no matter the type of nn.Module, that the order of execution of layers will stay the same? Even in weird net and hardware architectures (for example: random layers, parallel heads or subnets, multiple GPU exec, ...)

Also, is all(layer.executed for layer in summ.summary_list) a correct proof that the net summary was correctly built? I mean by that that a full forward pass was made when analyzing the network.

Thank you in advance for helping me out on this one.