Loading pretrained weights into new model

L-Reichardt: <p>Hello,</p>
<p>i have made modifications to a model, expanding it with multiple attention mechanisms. 80% of the model is the same as the previous version. Is it possible, to take the pretrained weights of the previous version, and insert them where applicable?</p>
<p>Ive trained the new model for 1 epoch, saving the weights (checkpoint). This is my attempt at updating those weights with pretrained weights:</p>
<pre><code>    checkpoint = torch.load(args.weight_path)
    pretrained = torch.load("pretrained_weights.pth")
    
    pretrained = {k: v for k, v in pretrained.items() if k in checkpoint}
    checkpoint.update(pretrained)
    
    G_mod.load_state_dict(checkpoint['state_dict'])</code></pre>

ptrblck: <p>Manipulating the <code>state_dict</code> sounds like a valid approach, if you’ve changed some parameters (e.g. in their shape).<br>
On the other hand, if you’ve added or removed parameters, you could try to use <code>load_state_dict(sd, strict=False)</code> to ignore missing or unexpected keys.</p>

L-Reichardt: <p>Thanks a lot<br>
This is exactly what I had in mind</p>

R_S: <p>hello again,</p>
<p>Can you please help me with how to use pretrained weights for different model. I have built a simple fcn model for segmentation with 5 blocks of conv, relu and pooling. I want to use pretrained weights of vgg/resent. I loaded the resnet pretrained weights as below. Is it the right way? I got better results, but I am not sure how the pretrained weights get added to my new model.</p>
<p>model = fcn()<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet50’]), strict=False)</p>
<p>Is it based on layer names or automatically concat the weights with corresponding filter sizes. Please give me some inputs, since I am new to using pretrained weights.</p>
<p>Thank you</p>

ptrblck: <p>I would not recommend to use <code>strict=False</code> without a proper verification that this workaround loads indeed the proper parameters.<br>
If you’ve changed the model architecture, try to load the parameters for each layer manually e.g. with a custom mapping which would load the parameters of the pretrained layer to the new one.</p>

R_S: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> okay, which is the good way?</p>
<ol>
<li>split the layers of pretrained model and concatinate with the layers of my custom model (according the layers output size which are required for fcn), then train the custom model.</li>
<li>use pretrained weights as features (remove final layers which are not required  and custom classifier layers and then train.<br>
for example in the second method i used vgg features,</li>
</ol>
<p>class fcn(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(fcn, self).<strong>init</strong>()<br>
self.features = vgg16.features<br>
self.classifier = nn.Sequential(nn.Conv2d(512, 4096, 7),nn.ReLU(inplace=True),nn.Conv2d(4096, 4096, 1),nn.ReLU(inplace=True),nn.Conv2d(4096, 32, 1),nn.ConvTranspose2d(32, 32, 224, stride=32)<br>
)</p>
<p>def forward(self, x):<br>
x = self.features(x)<br>
x = self.classifier(x)<br>
<span class="hashtag">#print</span>(x.shape)<br>
return x</p>
<p>Also, in both case, what is the affect of freezing the features. Should i train it or not?<br>
for param in vgg16.features.parameters():<br>
param.requires_grad = False</p>
<p>Thank you</p>

ptrblck: <ol>
<li>I don’t know how you could like to concatenate the layers and how this would be related to loading the pretrained weights, so could you explain this idea a bit more, please?</li>
<li>Yes, this approach looks valid.</li>
</ol>
<aside class="quote no-group" data-username="R_S" data-post="6" data-topic="104159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/r_s/48/42247_2.png" class="avatar"> R_S:</div>
<blockquote>
<p>Also, in both case, what is the affect of freezing the features. Should i train it or not?</p>
</blockquote>
</aside>
<p>It depends on your use case and you would have to run experiments to determine if freezing the layers would work for you. Often the feature extractor is frozen, if the new dataset is “similar” to the pretrained one (ImageNet), but I wouldn’t claim that’s a general rule.</p>

R_S: <p>for example this is my custom fcn model (only fcn 32)of 5 stages. As we have 5 pool layers, i want to add this pool output and pool output of vgg pretrained model. And then train this model. Is it possible?</p>
<p>class fcn(nn.Module):</p>
<pre><code>def __init__(self):
    super(fcn, self).__init__()
    self.conv1_1 = nn.Conv2d(3, 64, 3, 1, 1)
    #self.bn1_1 = nn.BatchNorm2d()
    self.relu1_1 = nn.ReLU()
    self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 1)
    #self.bn1_2 = nn.BatchNorm2d()
    self.relu1_2 = nn.ReLU()
    self.pool1 = nn.MaxPool2d(2, 2)#/2

    self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 1)
    # self.bn2_1 = nn.BatchNorm2d()
    self.relu2_1 = nn.ReLU()
    self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 1)
    # self.bn2_2 = nn.BatchNorm2d()
    self.relu2_2 = nn.ReLU()
    self.pool2 = nn.MaxPool2d(2, 2)#/4

    self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 1)
    # self.bn3_1 = nn.BatchNorm2d()
    self.relu3_1 = nn.ReLU()
    self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 1)
    # self.bn3_2 = nn.BatchNorm2d()
    self.relu3_2 = nn.ReLU()
    self.pool3 = nn.MaxPool2d(2, 2)#/8

    self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 1)
    # self.bn4_1 = nn.BatchNorm2d()
    self.relu4_1 = nn.ReLU()
    self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 1)
    # self.bn4_2 = nn.BatchNorm2d()
    self.relu4_2 = nn.ReLU()
    self.pool4 = nn.MaxPool2d(2, 2)  # /16

    self.conv5_1 = nn.Conv2d(512, 1024, 3, 1, 1)
    # self.bn5_1 = nn.BatchNorm2d()
    self.relu5_1 = nn.ReLU()
    self.conv5_2 = nn.Conv2d(1024, 1024, 3, 1, 1)
    # self.bn5_2 = nn.BatchNorm2d()
    self.relu5_2 = nn.ReLU()
    self.pool5 = nn.MaxPool2d(2, 2)  # /32

    self.fc1 = nn.Conv2d(1024, 4096, 3, 1, 1)
    self.relu6 = nn.ReLU()
    self.drop6 = nn.Dropout2d()

    self.fc2 = nn.Conv2d(4096, 4096, 1)
    self.relu7 = nn.ReLU()
    self.drop7 = nn.Dropout2d()

    self.fc3 = nn.Conv2d(4096, 32, 1)
    self.up = nn.ConvTranspose2d(32, 32, kernel_size=32, stride=32)

def forward(self, input):
    #print(input.shape)
    x = self.conv1_1(input)
    x = self.relu1_1(x)
    x = self.conv1_2(x)
    x = self.relu1_2(x)
    x = self.pool1(x)
    #print(x.shape)
    x = self.conv2_1(x)
    x = self.relu2_1(x)
    x = self.conv2_2(x)
    x = self.relu2_2(x)
    x = self.pool2(x)
    #print(x.shape)
    x = self.conv3_1(x)
    x = self.relu3_1(x)
    x = self.conv3_2(x)
    x = self.relu3_2(x)
    x = self.pool3(x)
    #print(x.shape)
    x = self.conv4_1(x)
    x = self.relu4_1(x)
    x = self.conv4_2(x)
    x = self.relu4_2(x)
    x = self.pool4(x)
    #print(x.shape)
    x = self.conv5_1(x)
    x = self.relu5_1(x)
    x = self.conv5_2(x)
    x = self.relu5_2(x)
    x = self.pool5(x)
    #print(x.shape)
    x = self.fc1(x)
    x = self.relu6(x)
    x = self.drop6(x)
    #print(x.shape)
    x = self.fc2(x)
    x = self.relu7(x)
    x = self.drop7(x)
    #print(x.shape)
    x = self.fc3(x)
    #print(x.shape)
    x = self.up(x)
    #print(x.shape)
    return x</code></pre>

ptrblck: <p>I don’t know exactly which layers are new, but since the model architecture was changed, you could create a mapping between the old and new layer names and load each layer separately.<br>
E.g something like this should work:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)

class MyNewModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.new_layer = nn.Linear(1, 1)

model = MyModel()
new_model = MyNewModel()

# this fails
new_model.load_state_dict(model.state_dict())
# &gt; RuntimeError: Error(s) in loading state_dict for MyNewModel:
#    Missing key(s) in state_dict: "new_layer.weight", "new_layer.bias". 
      
# this works, but could be dangerous, if you are not careful
new_model.load_state_dict(model.state_dict(), strict=False)
# &gt; _IncompatibleKeys(missing_keys=['new_layer.weight', 'new_layer.bias'], unexpected_keys=[])

# check
print((new_model.fc1.weight == model.fc1.weight).all())
# &gt; tensor(True)

# create mapping
mapping = [['fc1', 'fc1']]
for m in mapping:
    print('loading {} to {}'.format(m[1], m[0]))
    getattr(new_model, m[0]).load_state_dict(getattr(model, m[1]).state_dict())
</code></pre>

R_S: <p>thank you for your help. i want to express my understandings so far since i am a beginner, please tell me if i am wrong.</p>
<ol>
<li>i cannot use pretrained weights on new image dimension directly since they are trained on 224 images.</li>
<li>but we can use them if we extract the layers and add them into out new model. like this</li>
</ol>
<p>class fcn(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(fcn, self).<strong>init</strong>()<br>
self.features = vgg16.features<br>
self.classifier = nn.Sequential(nn.Conv2d(512, 4096, 7),nn.ReLU(inplace=True),nn.Conv2d(4096, 4096, 1),nn.ReLU(inplace=True),nn.Conv2d(4096, 32, 1),nn.ConvTranspose2d(32, 32, 224, stride=32)<br>
)</p>
<p>def forward(self, x):<br>
x = self.features(x)<br>
x = self.classifier(x)<br>
<span class="hashtag">#print</span>(x.shape)<br>
return x</p>
<p>why doesnt it give error? am i just using it as features?<br>
thank you</p>

ptrblck: <p>I don’t think these points are true, as you should be able to fine-tune the pretrained model using other input shapes as long as they are large enough (otherwise an intermediate activation could be reduced to an empty tensor and you would see an error).</p>

R_S: <p>thank you. I understand, we should check the output size so that they dont vanish</p>

R_S: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> hello, i have some issue with implementing fcn 32/16/8. I am using vgg16 pretrained weights and adding to my fcn model. For some reason my fcn 16 and 8 variations give bad results than fcn 32. Are these two commands same?</p>
<p>combined = pool4_scored + pool5_upscored<br>
combined = torch.cat([pool4_scored, pool5_upscored])</p>
<p>thank you</p>

ptrblck: <p>No, these operations are not doing the same, as adding tensors via:</p>
<pre><code class="lang-python">combined = pool4_scored + pool5_upscored
</code></pre>
<p>will perform the addition in each element while <code>torch.cat</code> will concatenate both tensors into a new (larger) tensor:</p>
<pre><code class="lang-python">a = torch.ones(3)
b = torch.ones(3) * 2

combined = a + b
print(combined)
&gt; tensor([3., 3., 3.])

combined = torch.cat((a, b))
print(combined)
&gt; tensor([1., 1., 1., 2., 2., 2.])
</code></pre>

R_S: <p>Thank you, can you please tell me which operation should I use for fcn architecture. I have used addition. I upscaled the vgg pool5 layers by 32 for fcn32 and gives good results. But when I add pool4 layers and (pool 5 up scale by 2) to get fcn16 output, it performs bad. According to the fcn architecture, fcn16 and fcn8 are supposed to work better.<br>
this is my code,<br>
import torch<br>
import torch.nn as nn<br>
import torchvision.models as models<br>
from pytorch_model_summary import summary</p>
<p>vgg16 = models.vgg16(pretrained=True)<br>
for param in vgg16.features.parameters():<br>
param.requires_grad = False<br>
<span class="hashtag">#False</span> Total params: 185,771,904 Trainable params: 171,057,216 Non-trainable params: 14,714,688<br>
<span class="hashtag">#true</span>  Total params: 185,771,904 Trainable params: 185,771,904 Non-trainable params: 0</p>
<p>class fcn(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(fcn, self).<strong>init</strong>()<br>
self.features = vgg16.features<br>
self.classifier = nn.Sequential(<br>
nn.Conv2d(512, 4096, 7),<br>
nn.ReLU(inplace=True),<br>
<span class="hashtag">#nn</span>.Dropout2d(),<br>
nn.Conv2d(4096, 4096, 1),<br>
nn.ReLU(inplace=True),<br>
<span class="hashtag">#nn</span>.Dropout2d(),<br>
nn.Conv2d(4096, 32, 1),<br>
nn.ConvTranspose2d(32, 32, 224, stride=32)<br>
)</p>
<p>def forward(self, x):<br>
x = self.features(x)#/32<br>
x = self.classifier(x)<br>
<span class="hashtag">#print</span>(x.shape)<br>
return x</p>
<p>class fcn16(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(fcn16, self).<strong>init</strong>()<br>
self.features = vgg16.features<br>
self.classifier = nn.Sequential(<br>
nn.Conv2d(512, 4096, 7),<br>
nn.ReLU(inplace=True),<br>
nn.Conv2d(4096, 4096, 1),<br>
nn.ReLU(inplace=True),<br>
nn.Conv2d(4096, 32, 1)<br>
)<br>
self.score_pool4 = nn.Conv2d(512, 32, 1)<br>
self.upscore2 = nn.ConvTranspose2d(32, 32, 14, stride=2, bias=False)<br>
self.upscore16 = nn.ConvTranspose2d(32, 32, 16, stride=16, bias=False)</p>
<p>def forward(self, x):<br>
pool4 = self.features<a>:-7</a><span class="hashtag">#512</span> features /16<br>
pool5 = self.features<a>-7:</a><span class="hashtag">#512</span> features /16/2=/32<br>
pool5_upscored = self.upscore2(self.classifier(pool5))<span class="hashtag">#32</span> class features stride2 /32*2=/16<br>
pool4_scored = self.score_pool4(pool4)<span class="hashtag">#32</span> features /16<br>
combined = pool4_scored + pool5_upscored<br>
<span class="hashtag">#combined</span> = torch.cat([pool4_scored, pool5_upscored])<br>
res = self.upscore16(combined)# /1<br>
return res</p>
<p>class fcn8(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(fcn8, self).<strong>init</strong>()<br>
self.features = vgg16.features<br>
self.classifier = nn.Sequential(<br>
nn.Conv2d(512, 4096, 7),<br>
nn.ReLU(inplace=True),<br>
nn.Conv2d(4096, 4096, 1),<br>
nn.ReLU(inplace=True),<br>
nn.Conv2d(4096, 32, 1)<br>
)<br>
self.score_pool4 = nn.Conv2d(512, 32, 1)<br>
self.score_pool3 = nn.Conv2d(256, 32, 1)<br>
self.upscore2 = nn.ConvTranspose2d(32, 32, 14, stride=2, bias=False)<br>
self.upscore3 = nn.ConvTranspose2d(32, 32, 2, stride=2, bias=False)<br>
<span class="hashtag">#self</span>.upscore16 = nn.ConvTranspose2d(32, 32, 16, stride=16, bias=False)<br>
self.upscore8 = nn.ConvTranspose2d(32, 32, 8, stride=8, bias=False)</p>
<p>def forward(self, x):<br>
pool3 = self.features<a>:-14</a><span class="hashtag">#256</span> features /8<br>
pool4 = self.features<a>-14:-7</a><span class="hashtag">#512</span> features /8/2=16<br>
pool5 = self.features<a>-7:</a><span class="hashtag">#512</span> features /16/2=/32<br>
pool5_upscored = self.upscore2(self.classifier(pool5))<span class="hashtag">#32</span> class features stride2 /32<em>2=/16<br>
pool4_scored = self.score_pool4(pool4)<span class="hashtag">#32</span> class features /16<br>
pool3_scored = self.score_pool3(pool3)<span class="hashtag">#32</span> class features /8<br>
combined = pool4_scored + pool5_upscored #/16<br>
<span class="hashtag">#print</span>(combined.shape)<br>
combined_upscored = self.upscore3(combined)<span class="hashtag">#32</span> class features stride2 /16</em>2=/8<br>
<span class="hashtag">#print</span>(combined_upscored.shape)<br>
combined2 = pool3_scored + combined_upscored<br>
<span class="hashtag">#print</span>(combined2.shape)<br>
<span class="hashtag">#res</span> = self.upscore16(combined)#/1<br>
res = self.upscore8(combined2)#/1<br>
<span class="hashtag">#print</span>(res.shape)<br>
return res</p>
<p>Thank you</p>

R_S: <p>i tried this way of taking weights from pretrained model. But this gives slightly bad accuracy than the original model.</p>
<pre><code>def forward(self, input):
    #print(input.shape)
    self.conv1_1.weight = vgg16.features[0].weight
    self.conv1_1.bias = vgg16.features[0].bias
    x = self.conv1_1(input)
    x = self.relu1_1(x)
    self.conv1_2.weight = vgg16.features[2].weight
    self.conv1_2.bias = vgg16.features[2].bias
    x = self.conv1_2(x)
    x = self.relu1_2(x)
    x = self.pool1(x)
    #print(x.shape)
    self.conv2_1.weight = vgg16.features[5].weight
    self.conv2_1.bias = vgg16.features[5].bias
    x = self.conv2_1(x)
    x = self.relu2_1(x)
    self.conv2_2.weight = vgg16.features[7].weight
    self.conv2_2.bias = vgg16.features[7].bias
    x = self.conv2_2(x)
    x = self.relu2_2(x)
    x = self.pool2(x)
    #print(x.shape)
    self.conv3_1.weight = vgg16.features[10].weight
    self.conv3_1.bias = vgg16.features[10].bias
    x = self.conv3_1(x)
    x = self.relu3_1(x)
    self.conv3_2.weight = vgg16.features[12].weight
    self.conv3_2.bias = vgg16.features[12].bias
    x = self.conv3_2(x)
    x = self.relu3_2(x)
    x = self.pool3(x)
    #print(x.shape)
    self.conv4_1.weight = vgg16.features[17].weight
    self.conv4_1.bias = vgg16.features[17].bias
    x = self.conv4_1(x)
    x = self.relu4_1(x)
    self.conv4_2.weight = vgg16.features[19].weight
    self.conv4_2.bias = vgg16.features[19].bias
    x = self.conv4_2(x)
    x = self.relu4_2(x)
    x = self.pool4(x)
    #print(x.shape)
    x = self.conv5_1(x)
    x = self.relu5_1(x)
    x = self.conv5_2(x)
    x = self.relu5_2(x)
    x = self.pool5(x)
    #print(x.shape)
    x = self.fc1(x)
    x = self.relu6(x)
    x = self.drop6(x)
    #print(x.shape)
    x = self.fc2(x)
    x = self.relu7(x)
    x = self.drop7(x)
    #print(x.shape)
    x = self.fc3(x)
    #print(x.shape)
    x = self.up(x)
    #print(x.shape)
    return x</code></pre>

ptrblck: <p>I’m not sure if your current use case changed, but you should not reassign parameters in the forward method, as they won’t be trained (in case you are fine-tuning the model).<br>
Did you try to use the previously posted approach?</p>

R_S: <p>previous method works good for fcn 32. but bad for fcn16 andf fcn8.</p>

ptrblck: <p>Could you explain a bit what “works bad” would mean in this context?<br>
Is the parameter loading not working? If so, could you post a code snippet showing the errors?</p>

R_S: <p>this is my code. I get less accuracy for fcn16 and 8, fcn32 gives good accuracy.</p>
<pre><code class="lang-auto">class fcn8(nn.Module):
def **init**(self):
super(fcn8, self).**init**()
self.features = vgg16.features
self.classifier = nn.Sequential(
nn.Conv2d(512, 4096, 7),
nn.ReLU(inplace=True),
nn.Conv2d(4096, 4096, 1),
nn.ReLU(inplace=True),
nn.Conv2d(4096, 32, 1)
)
self.score_pool4 = nn.Conv2d(512, 32, 1)
self.score_pool3 = nn.Conv2d(256, 32, 1)
self.upscore2 = nn.ConvTranspose2d(32, 32, 14, stride=2, bias=False)
self.upscore3 = nn.ConvTranspose2d(32, 32, 2, stride=2, bias=False)
#self.upscore16 = nn.ConvTranspose2d(32, 32, 16, stride=16, bias=False)
self.upscore8 = nn.ConvTranspose2d(32, 32, 8, stride=8, bias=False)

def forward(self, x):
pool3 = self.features:-14#256 features /8
pool4 = self.features-14:-7#512 features /8/2=16
pool5 = self.features-7:#512 features /16/2=/32
pool5_upscored = self.upscore2(self.classifier(pool5))#32 class features stride2 /32&lt;em&gt;2=/16
pool4_scored = self.score_pool4(pool4)#32 class features /16
pool3_scored = self.score_pool3(pool3)#32 class features /8
combined = pool4_scored + pool5_upscored #/16
#print(combined.shape)
combined_upscored = self.upscore3(combined)#32 class features stride2 /16&lt;/em&gt;2=/8
#print(combined_upscored.shape)
combined2 = pool3_scored + combined_upscored
#print(combined2.shape)
#res = self.upscore16(combined)#/1
res = self.upscore8(combined2)#/1
#print(res.shape)
return res
</code></pre>