Transfer learning from one network architecture to another network architecture

maryma: <p>Hi. I have one model that is efficientnet it has some layers and one average pooling layer and last layers are some fully connected layers. This is for image classification you can assume we have 5 classes here. I trained the model with 1000, 3d MRI images and I saved the model by load.state_dict. Then I have another image classification model that it has 10 classes and I want to train it by 200, 3d images. But I wanted to use the first model weights for the second model before the fully connected layers. How is it possible?<br>
I know until here<br>
I have to make an object of first classification<br>
<code>Classification1 = Classfication1_Model()</code><br>
Load torch weights<br>
<code>Classification1.load_state_dict(torch.load(path_to_state_dict))</code><br>
I have to make an object of second classification<br>
But Here it is the exact question I have. How can I use pretrained model weights for new classification? Should I freeze layers before average pooling layer and how can I add weights to model?</p>

azhanmohammed: <p>I am assuming that the network architecture in both models is exactly the same except for the classification layer. In that case, you can load the weights using the <code>load_state_dict</code> function, and then use <code>Classification.linear=nn.Linear(input_feature,200)</code> to change the classification layer and fine-tune using the previous weights you got.</p>

maryma: <p>How can I fill input_feature? There is no need to fill it? and instead of 200 you mean 10? Beacuse in second I have 10 classes.<br>
But I have one important question. If I had one segmentation model and wanted to use the weights of that model for classification model what should I do? How can I use the weights in this situation?</p>

azhanmohammed: <aside class="quote no-group" data-username="maryma" data-post="3" data-topic="159575">
<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/maryma/48/56556_2.png" class="avatar"> maryma:</div>
<blockquote>
<p>and instead of 200 you mean 10? Beacuse in second I have 10 classes.</p>
</blockquote>
</aside>
<p>Yes, basically replace the last output of linear layer with your new number of classes.</p>
<aside class="quote no-group" data-username="maryma" data-post="3" data-topic="159575">
<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/maryma/48/56556_2.png" class="avatar"> maryma:</div>
<blockquote>
<p>If I had one segmentation model and wanted to use the weights of that model for classification model what should I do? How can I use the weights in this situation?</p>
</blockquote>
</aside>
<p>A segmentation model has an encoder decoder architecture, you can use the weights of the encoder part, add a linear layer on top for classification and use it for classification as well.</p>

maryma: <p>I am familiar with the theory of it but I don’t know how should I apply it?<br>
Can you please explain it with code?<br>
My questions are exactly these ones:<br>
1)Because when we save a model with pth extension all weights of model were saved but I need just the encoder part weights how can I obtain it?<br>
2)How can I make a classification model that has only the encoder part of segmentation(for example Unet3d) and add a layer of classification on top of it?</p>

azhanmohammed: <p><a class="mention" href="/u/maryma">@maryma</a>  sure, I can provide a minimal code that can help you in understanding how to use the saved model.</p>
<aside class="quote no-group" data-username="maryma" data-post="5" data-topic="159575">
<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/maryma/48/56556_2.png" class="avatar"> maryma:</div>
<blockquote>
<p>1)Because when we save a model with pth extension all weights of model were saved but I need just the encoder part weights how can I obtain it?</p>
</blockquote>
</aside>
<p>We first define the auto encoder model which does both encoding and decoding:</p>
<pre data-code-wrap="python"><code class="lang-python">class encoderDecoder(nn.Module):
    def __init__(self):
        super(encoderDecoder, self).__init__()
        self.encoder = nn.Sequential(#encoder layers go here)
        self.decoder = nn.Sequential(#decoder layers go here)
        
    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x
</code></pre>
<p>We can save its weight using:</p>
<pre data-code-wrap="python"><code class="lang-python">model = encoderDecoder()
torch.save(model.state_dict(), 'autoencoder.pt')
</code></pre>
<p>Next we define a model but with only encoder layers:</p>
<pre data-code-wrap="python"><code class="lang-python">class encoderModel(nn.Module):
    def __init__(self):
        super(encoderModel, self).__init__()
        self.encoder = nn.Sequential(nn.Conv2d(3,16,3))
        
    def forward(self, x):
        x = self.encoder(x)
        return x
</code></pre>
<p>Now we can load the weights using:</p>
<pre data-code-wrap="python"><code class="lang-python">model2 = encoderModel()
model2.load_state_dict(torch.load('autoencoder.pt'), strict=False)
</code></pre>
<p>You will get a warning that says:</p>
<pre data-code-wrap="python"><code class="lang-python">_IncompatibleKeys(missing_keys=[], unexpected_keys=['decoder.0.weight', 'decoder.0.bias'])
</code></pre>
<p>But it should be fine since we do not need the decoder layers here. So this clears how we can just use the encoder part for our model.</p>
<aside class="quote no-group" data-username="maryma" data-post="5" data-topic="159575">
<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/maryma/48/56556_2.png" class="avatar"> maryma:</div>
<blockquote>
<p>2)How can I make a classification model that has only the encoder part of segmentation(for example Unet3d) and add a layer of classification on top of it?</p>
</blockquote>
</aside>
<p>For this part we will use the same encoder decoder model that we used above:</p>
<pre data-code-wrap="python"><code class="lang-python">class encoderDecoder(nn.Module):
    def __init__(self):
        super(encoderDecoder, self).__init__()
        self.encoder = nn.Sequential(#encoder layers go here)
        self.decoder = nn.Sequential(#decoder layers go here)
        
    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x
</code></pre>
<p>We save this model in the same way:</p>
<pre data-code-wrap="python"><code class="lang-python">model = encoderDecoder()
torch.save(model.state_dict(), 'autoencoder.pt')
</code></pre>
<p>For the classification part we need to add a linear layer (or few layers if you want). This can be done as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">class classificationModel(nn.Module):
    def __init__(self):
        super(classificationModel, self).__init__()
        self.encoder = nn.Sequential(#encoder layers go here)
        self.fc = nn.Sequential(nn.Flatten(), nn.Linear(#add input number of features here, #add number of classes here))
        
    def forward(self, x):
        x = self.encoder(x)
        x = self.fc(x)
        return x
</code></pre>
<p>And you can load weights just as we did before using:</p>
<pre data-code-wrap="python"><code class="lang-python">model2 = classificationModel()
model2.load_state_dict(torch.load('autoencoder.pt'), strict=False)
</code></pre>
<p>This time you will get another warning which says that it cannot find state dict keys for fc layer, but that is fine as the original model did not have those layers. This way you can use the encoder part to train a classifier.</p>

maryma: <p>Thanks for your complete answere</p>