How to find input layers names for intermediate layer in PyTorch model?

Question

I have some complicated model on PyTorch. How can I print names of layers (or IDs) which connected to layer's input. For start I want to find it for Concat layer. See example code below:

class Concat(nn.Module):
    def __init__(self, dimension=1):
        super().__init__()
        self.d = dimension

    def forward(self, x):
        return torch.cat(x, self.d)


class SomeModel(nn.Module):
    def __init__(self):
        super(SomeModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        self.conc = Concat(1)
        self.linear = nn.Linear(8192, 1)

    def forward(self, x):
        out1 = F.relu(self.bn1(self.conv1(x)))
        out2 = F.relu(self.conv2(x))
        out = self.conc([out1, out2])
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


if __name__ == '__main__':
    model = SomeModel()
    print(model)
    y = model(torch.randn(1, 3, 32, 32))
    print(y.size())
    for name, m in model.named_modules():
        if 'Concat' in m.__class__.__name__:
            print(name, m, m.__class__.__name__)
            # Here print names of all input layers for Concat

Ivan · Accepted Answer · 2021-09-29 21:36:02Z

2

+50

You can use type(module).__name__ to get the nn.Module class name:

>>> model = SomeModel()
>>> y = model(torch.randn(1, 3, 32, 32))
>>> for name, m in model.named_modules():
...     if 'Concat' == type(m).__name__:
...         print(name, m)
conc Concat()

Edit: You can actually manage to get the list of operators used to compute the inputs of Concat. However, I don't think you can actually get the attribute names of the nn.Module associated with these operators. This kind of information is not available - and needed - at model inference.

This solution requires you to register a forward hook on the layer with nn.Module.register_forward_hook. Then perform one inference to trigger it, then you can remove the hook. In the forward hook, you have access to the list of inputs and extract the name of the operator from the grad_fn attribute callback. Using nn.Module.register_forward_pre_hook here would be more appropriate since we are only looking at the inputs, and do not need the output.

>>> def op_name(x)
...     return type(x.grad_fn).__name__.replace('Backward0', '')

>>> def forward_hook(module, ins):
...     print([op_name(x) for x in ins[0]])

Attach the hook on model.conc, trigger it and then clean up:

>>> handle = model.conc.register_forward_pre_hook(forward_hook)
>>> model(torch.empty(2, 3, 10, 10, requires_grad=True))
['Relu', 'Relu']

>>> handle.remove()

edited Sep 29, 2021 at 21:36

answered Sep 24, 2021 at 16:52

Ivan

41.3k9 gold badges78 silver badges120 bronze badges

Sign up to request clarification or add additional context in comments.

7 Comments

ZFTurbo Over a year ago

I need name of inputs for this layer. I know how to print name of current layer.

Ivan Over a year ago

@ZFTurbo What do you mean by "name"? What would be the expected output in your example?

ZFTurbo Over a year ago

Example: ['module.conv1', 'module.conv2']

Ivan Over a year ago

@ZFTurbo why is the expected result not ['module.bn1', 'module.conv2'], since bn1 is applied after conv1 before being passed to Concat? Do you want to return the previous closests conv2d layer that were used on each input of Concat?

Ivan Over a year ago

@ZFTurbo, I have edited my answer.

|

MWB · Accepted Answer · 2021-09-24 19:29:46Z

1

How can I print names of layers (or IDs) which connected to layer's input.

You can't do it based on a module itself, because "connected to" doesn't exist as a static property of a nn.Module:

Instead of

out = self.conc([out1, out2])

you could just as easily have written

out = self.conc([out1, random.choice([out1, out2]))

What's connected to self.conc now?

However, you can inspect the "graph" of a value (If torchviz.make_dot doesn't directly do what you want, you may want to check its implementation)

answered Sep 24, 2021 at 19:29

MWB

12.7k9 gold badges55 silver badges107 bronze badges

Collectives™ on Stack Overflow

How to find input layers names for intermediate layer in PyTorch model?

2 Answers 2

7 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

7 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related