Editing layers of pretrained model make finetuning fail

username: jchaykow
cooked: <p>If I overwrite a layer in the pretrained model the training is fine, but if I remove the final layer and then re-append the same layer - the training fails.</p>
<p>I’m starting here: <a href="https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html</a></p>
<p>Here’s my final code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.models as models
from torch.autograd import Variable

# Pretrained resnet152 model till second last layer as feature extraction.
resnet152 = models.resnet152(pretrained=True)

last_layer = list(resnet152.children())[-1]
modules=list(resnet152.children())[:-1]
resnet152=nn.Sequential(*modules)
for p in resnet152.parameters():
    p.requires_grad = False

resnet152.fc = last_layer

optimizer_conv = optim.SGD(resnet152.fc.parameters(), lr=0.001, momentum=0.9)

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=5, gamma=0.1)

model_ft = train_model(resnet152, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=10)
</code></pre>
<p>This will fail with:</p>
<p><code>RuntimeError: size mismatch, m1: [8192 x 1], m2: [2048 x 1000] at ../aten/src/TH/generic/THTensorMath.cpp:961</code></p>
<p>But if I do the same thing but instead overwrite the .fc layer:</p>
<pre><code class="lang-auto">resnet152 = models.resnet152(pretrained=True)

for p in resnet152.parameters():
    p.requires_grad = False

num_ftrs = resnet152.fc.in_features
resnet152.fc = nn.Linear(num_ftrs, 1000)

optimizer_conv = optim.SGD(resnet152.fc.parameters(), lr=0.001, momentum=0.9)

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=5, gamma=0.1)

model_ft = train_model(resnet152, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=10)
</code></pre>
<p>This will train fine.</p>
<p>Even if I check they are equivalent it passes:</p>
<pre><code class="lang-auto">resnet152_2 = models.resnet152(pretrained=True)

for param, param2 in zip(resnet152.parameters(), resnet152_2.parameters()):
    assert torch.equal(param[0], param2[0])
</code></pre>

username: ptrblck
cooked: <p>If you wrap all modules in an <code>nn.Sequential</code> module, you are not using the original <code>forward</code> method anymore, and thus are missing <a href="https://github.com/pytorch/vision/blob/6693b2c6890a46734bbbfef9a3910db0d6b11c87/torchvision/models/resnet.py#L206" rel="nofollow noopener">this reshape operation</a>.</p>
<p>You could add a custom <code>Flatten</code> layer or write your own <code>forward</code> method using a custom <code>nn.Module</code>:</p>
<pre><code class="lang-python">class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        return x

last_layer = list(resnet152.children())[-1]
modules=list(resnet152.children())[:-1]
resnet152=nn.Sequential(*modules)
for p in resnet152.parameters():
    p.requires_grad = False

resnet152.fc = nn.Sequential(
    Flatten(),
    last_layer
)

x = torch.randn(2, 3, 224, 224)
output = resnet152(x)
</code></pre>

username: jchaykow
cooked: <p>Thank you! So I assume if I had a <code>F.normalize(x, p=2, dim=1)</code> in my forward pass I would just replace with:</p>
<pre><code class="lang-auto">class normalize(nn.Module):
    def __init__(self):
        super(normalize, self).__init__()
        
    def forward(self, x):
        x = F.normalize(x, p=2, dim=1)
        return x
</code></pre>

username: ptrblck
cooked: <p>Yes, all functional calls in the original <code>forward</code> will be lost, if you rewrap the submodules in a new module, so your approach should be correct.</p>