Why torchvision.models.alexnet() not work when i test the input and output?

he Code is here:

import torch
from torchvision.models import alexnet

if __name__ == "__main__":
    net = alexnet()
    x = torch.rand((1, 3, 224, 224))

    for name, layer in net.named_children():
        x = layer(x)
        print(name, ' output shape:\t', x.shape)

The output is here:

features  output shape:  torch.Size([1, 256, 6, 6])
avgpool  output shape:   torch.Size([1, 256, 6, 6])

File "AlexNet.py", line 9, in <module>
    x = layer(x)
RuntimeError: size mismatch, m1: [1536 x 6], m2: [9216 x 4096]

I want to test the input and output of the AlexNet, but i failed with the ‘size mismatch’ error. i use the image size(3,224,224) provided by the paper of AlexNet, and i wanna get the right output. After this error, i try to fix the

self.avgpool = nn.AdaptiveAvgPool2d((6, 6))

to

self.avgpool = nn.AdaptiveAvgPool2d((1, 9216))

and i get the right output.

I really want to know what did i do wrong with the torchvision. And i also want to know what people would do when they need to test the input and output of the cnn.

thx for your help!

You are missing the flatten operation in this line of code if you call all modules in a sequential way.

One approach is to use forward hooks to get the output without manipulating the forward method.

1 Like

Your answer is greatly appreciated, it solved my problem.Thank you very much! And i also want to ask, if I want to see the output size in the middle of the network, is there a better way? What I know so far is using net.named_children () get the ouput after layer by layer calculation, another way is manually calculate the size

One approach would be to create a new custom model class by deriving from alexnet as the base class and add print statements to the forward method.
However, this might be too cumbersome just to get the output shapes.

I usually just register hooks for all modules, which seems to work fine for some quick debugging:

model = models.alexnet()

for name, child in model.named_modules():
    if isinstance(child, nn.Module):
        child.register_forward_hook(
            lambda m, inp, out, name=name: print('{} output shape {}'.format(name, out.shape)))

x = torch.randn(1, 3, 224, 224)
output = model(x)

> features.0 output shape torch.Size([1, 64, 55, 55])
features.1 output shape torch.Size([1, 64, 55, 55])
features.2 output shape torch.Size([1, 64, 27, 27])
features.3 output shape torch.Size([1, 192, 27, 27])
features.4 output shape torch.Size([1, 192, 27, 27])
features.5 output shape torch.Size([1, 192, 13, 13])
features.6 output shape torch.Size([1, 384, 13, 13])
features.7 output shape torch.Size([1, 384, 13, 13])
features.8 output shape torch.Size([1, 256, 13, 13])
features.9 output shape torch.Size([1, 256, 13, 13])
features.10 output shape torch.Size([1, 256, 13, 13])
features.11 output shape torch.Size([1, 256, 13, 13])
features.12 output shape torch.Size([1, 256, 6, 6])
features output shape torch.Size([1, 256, 6, 6])
avgpool output shape torch.Size([1, 256, 6, 6])
classifier.0 output shape torch.Size([1, 9216])
classifier.1 output shape torch.Size([1, 4096])
classifier.2 output shape torch.Size([1, 4096])
classifier.3 output shape torch.Size([1, 4096])
classifier.4 output shape torch.Size([1, 4096])
classifier.5 output shape torch.Size([1, 4096])
classifier.6 output shape torch.Size([1, 1000])
classifier output shape torch.Size([1, 1000])
 output shape torch.Size([1, 1000])

Thank you very much for sharing your experience. I learned a lot about debugging a cnn. :hugs:

1 Like