Accessing operations in `forward()`

1

Let’s say I have the following model

class Model(nn.Module):

    def __init__(self):
        super().__init__()

        self.conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=(3, 3))

    def forward(self, x):
        x = self.conv(x)
        x = torch.flatten(x, start_dim=1)
        x = torch.sigmoid(x)
        return x

Now, if I use

for name, module in model.named_children():
    print(type(module))

only the convolutional operations shows up:

<class 'torch.nn.modules.conv.Conv2d'>

I would like to know, why the flatten() and sigmoid() do not show up?

If I, for example, want to replace in an existing model the flatten() or sigmoid() operation, how do I do that in this case? How can I modify the forward pass?

2

The functional operations used in the forward method are not registered as modules and thus won’t show up in model.modules() or model.children().

You can directly modify the forward method and replace these operations.
In case you want to manipulate the forward of another pre-defined model, you could derive a custom model from it and implement your forward there.

3

Hello, take Samue1’s example, if I want to get the output tensor of unregistered operation torch.flatten in forward what can I do?

4

You could return this tensor in the forward additionally to the last output.

5

Thanks for your quick reply! Sorry, I didn’t make my condition clear, I cannot modify the forward function as I wish since my users don’t want me to. Is there any function or way to get output tensors of desired ops by name or something like in tensorflow1?

6

I don’t think this is easily doable. Forward hooks need nn.Modules, while it seems you are using the pure functional API in the forward. I don’t know if e.g. torch.fx would be able to manipulate these functional calls to modules, which you could then use in forward hooks, but it might be worth checking the docs.

7

Will give it a try, thanks!