为什么在测试输入和输出时torchvision.models.alexnet（）不起作用？

Question

代码在这里：

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)

输出在这里：

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]

我想测试AlexNet的输入和输出，但是由于“大小不匹配”错误而失败。我使用AlexNet论文提供的图像大小（3,224,224），我想获得正确的输出。发生此错误后，我尝试修复

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

到

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

我得到正确的输出。

我真的很想知道我对火炬手做了什么错。我也想知道人们在需要测试cnn的输入和输出时会怎么做。

谢谢您的帮助！

Answer 1

我已经解决了这个问题。我通过这种方式错过了展平的x = torch.flatten(x, 1) code

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

谢谢您的帮助！