How to obtain the shape of the output tensor from each step in a Pytorch pretrained model

So I am trying to use the Pytorch implementation of the VGG16 model. In other words, I downloaded the model and waits through the usual Pytorch API.

I need to make some changes at the bottom of the pre-trained model, to accommodate my own data. I was wondering if there is a way to “step through” the model with some random data, so that I can check the sizes of the input and output tensors for each layer? Now I could do a print(model) and see the summary information for each layer, and that is fine. But the summary info does not always match up exactly with the tensors coming into that layer–because the layer will automatically adjust itself for different batch sizes or other settings.

Hence, I wanted to see if I could just generate some data and step through a pretrained model. Is there a way to do this, using pdb or something? I guess part of my confusion is I could not find in the documentation the structure of these model objects.

Thanks.

You could step into the forward pass like this:

x = torch.randn((1,3,28,28))
y = pretrained_model(x)

Make sure you step into y = pretrained_model(x) using pdb. Eventually, you will reach the forward method of the pre-trained model’s class, after which you can probe the different tensor shapes.

Also, see the VGG class’s forward method here.

Hope this helps.

@Mahmoud_Abdelkhalek Ahh yes, that make sense. I get you now. That is great, I can just step into the model up till the forward() method, and then track from there. Thanks so much. Haha, that was easier than I thought. I was probably overthinking it, figuring that it would be harder than that. But what you suggested is easy enough. Thanks again :).

from torchinfo import summary
summary(model_0, input_size=[1, 3, 64, 64])

This outputs a very nice summary of all that is happening in your model.