Here is a simple example. I tried to divide a network (Resnet50) into two parts: head
and tail
using children
. Conceptually, this should work but it doesn't. Why is it?
import torch
import torch.nn as nn
from torchvision.models import resnet50
head = nn.Sequential(*list(resnet.children())[:-2])
tail = nn.Sequential(*list(resnet.children())[-2:])
x = torch.zeros(1, 3, 160, 160)
resnet(x).shape # torch.Size([1, 1000])
head(x).shape # torch.Size([1, 2048, 5, 5])
tail(head(x)).shape # Error: RuntimeError: size mismatch, m1: [2048 x 1], m2: [2048 x 1000] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:136
For information, the tail is nothing but
Sequential(
(0): AdaptiveAvgPool2d(output_size=(1, 1))
(1): Linear(in_features=2048, out_features=1000, bias=True)
)
So I actually know that if I can do like this. But then, why the reshaping function (view
) is not in the children?
pool =resnet._modules['avgpool']
fc = resnet._modules['fc']
fc(pool(head(x)).view(1, -1))
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…