Is there an inbuilt function to print the output of each neuron of a CNN in PyTorch

Question

I created one CNN model with 4 convolutional layers. how to find the number of neurons and output of the each layer

[CNN Architecture ]

(https://i.stack.imgur.com/XASaf.png)```

Answer 1

To give a short answer to your question: No.

---

You can create a module inspector by attaching hooks. An easy method is to prepare an empty dictionary before inference and feed the results via the hook calls.

Here is a minimal example using nn.Module.register_forward_hook.

Construct a hook factory function:

def create_hook(name, out):
    def hook(module, in_tensor, out_tensor):
        out[name] = get_parameter_count(module), out_tensor
    return hook

Now you can attach the custom hooks on the different layers of the module named model:

out = {}
for name, layer in model.named_children():
    layer.register_forward_hook(create_hook(name, out))

Finally we are ready to infer by calling model. The intermediate results and submodule sizes will be contained inside out after inference.

---

Here is a demo with the following class:

class DummyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.convs = nn.Sequential(nn.Conv2d(3, 16, 3, 3), 
                                   nn.Conv2d(16, 8, 3, 3))
        self.pool = nn.AdaptiveAvgPool2d(2)
        self.classif = nn.LazyLinear(10)

    def forward(self, x):
        out = self.convs(x)
        out = self.pool(out)
        out = self.classif(out)
        return out

And the desired dictionary is constructed with the following:

>>> model = DummyNet()
>>> y = model(torch.rand(2, 3, 100, 100))

>>> for name, (c, y) in out.items():
...     print(f'[{name}] #parameter={c}, output shape={y[0].shape}')
[convs] #parameter=1608, output shape=torch.Size([8, 11, 11])
[pool] #parameter=0, output shape=torch.Size([8, 2, 2])
[classif] #parameter=30, output shape=torch.Size([8, 2, 10])

You can read more about forward hooks in another thread: How to get all the tensors in a graph?