How to get the input and output feature maps of a CNN model?

1

Hi, I am trying to find the dimensions of an image as it goes through a convolutional neural network at each layer. So for instance, if there is maxpooling or convolution being applied, I’d like to know the shape of the image at that layer, for all layers. I know I can use the nOut=image+2p-f / s + 1 formula but it would be too tedious and complex given the size of the model. Is there a simple way to do this, perhaps a visualization tool/script or something? Thanks in advance.

2

I’ve never implemented this before, but I think you can iterate through all the Modules in your model and register a hook that prints the shape of the output, using register_forward_hook like here.

Something roughly like this:

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)

def hook(module, input, output):
    print(output.shape)
    return None

for module in model.children():
    module.register_forward_hook(hook)
    
_ = model(torch.rand(1, 3, 224, 224))

Output:
torch.Size([1, 64, 112, 112])
torch.Size([1, 64, 112, 112])
torch.Size([1, 64, 112, 112])
torch.Size([1, 64, 56, 56])
torch.Size([1, 64, 56, 56])
torch.Size([1, 128, 28, 28])
torch.Size([1, 256, 14, 14])
torch.Size([1, 512, 7, 7])
torch.Size([1, 512, 1, 1])
torch.Size([1, 1000])

Note that in this case some of the children are Sequential blocks, so you may want to step through those / print their names and such.