Save and Load multi model

Praveen_K: <p>I’m trying to train <code>Autoencoder</code> as image classifier. Here is my code for <code>Autoencoder</code></p>
<pre><code class="lang-auto">class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        # encoder
        self.encoder = nn.Sequential(nn.Linear(in_features=784, out_features=16), 
                                  nn.ReLU()
                                 )
        # decoder
        self.decoder = nn.Sequential(nn.Linear(in_features=16, out_features=784), 
                                     nn.ReLU(),
                                    )

   
    def forward(self, x):
        x = self.encoder(x)       
        x = self.decoder(x)

        return x

net = Autoencoder()
</code></pre>
<p>After training of <code>Autoencoder</code> I’m training classifier</p>
<pre><code class="lang-auto"># copy encoder part
new_classifier = nn.Sequential(*list(net.children())[:-1])
net = new_classifier
# add FC and Softmax
net.add_module('classifier', nn.Sequential(nn.Linear(100, x),nn.LogSoftmax(dim=1)))
</code></pre>
<p>Now I would like to Save and Load model in the below format.</p>
<pre><code class="lang-auto">torch.save(the_model.state_dict(), PATH)
</code></pre>
<p>Then later:</p>
<pre><code class="lang-auto">the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
</code></pre>
<p>How can I load the model as above if I have 2 parts(<code>Encoder</code> and <code>Classifier</code>) ?<br>
It is little bit confusing in Multi models save and load in documentation. It would be helpful if someone helps me with this.</p>
<p>Thanks</p>

Praveen_K: <p>Can anyone pls help me with this ?</p>

ptrblck: <p>I’m not sure I understand the use case correctly.</p>
<aside class="quote no-group" data-username="Praveen_K" data-post="1" data-topic="111445">
<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/praveen_k/48/27741_2.png" class="avatar"> Praveen_K:</div>
<blockquote>
<p>How can I load the model as above if I have 2 parts(<code>Encoder</code> and <code>Classifier</code>) ?</p>
</blockquote>
</aside>
<p>You can directly save and load the <code>state_dict</code> of the <code>Autoencoder</code> even if it contains the two submodules.<br>
Wrapping the model into an <code>nn.Sequential</code> container might work for your use case, but would change the <code>state_dict</code> keys.<br>
Saving the <code>state_dict</code> of the <code>nn.Sequential</code> container and loading it into the <code>Autoencoder</code> will most likely not work out of the box and you would need to adapt the keys inside the <code>state_dict</code>.</p>

Praveen_K: <p>I understood about adapting keys. I tried few things myself.</p>
<p>Here is how I’m saving and loading the model.</p>
<pre><code class="lang-auto">
torch.save(
        {
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, './model.pth')


</code></pre>
<p>Loading the model…</p>
<pre><code class="lang-auto">class Initialise(nn.Module):
    def __init__(self):
        super(Initialise, self).__init__()
        self.encoder = nn.Sequential(nn.Linear(28*28*3, 256), nn.ReLU())
        self.classifier = nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1))
  
    def forward(self, x):
        x = self.encoder(x)
        x = self.classifier(x)

    
        return x

model= Initialise()
checkpoint = torch.load('./model.pth', map_location='cuda')
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
</code></pre>
<p>But at the moment I’m facing the below error.</p>
<p><code>_IncompatibleKeys(missing_keys=['encoder.0.weight', 'encoder.0.bias'], unexpected_keys=['0.0.weight', '0.0.bias'])</code></p>
<p>If I understand properly, I should just have to change layer name before saving from <code>0</code> to <code>encoder </code>, right?</p>

ptrblck: <p>Yes, the issue is raised, since you are changing the model definition from a custom module using the <code>self.encoder</code> and <code>self.classifier</code> submodules to a sequential module.<br>
One fix would be to rename the keys in the <code>state_dict</code>. However, I’m currently unsure, if you really need to change the model definition or if it would be possible to use the custom model definition.<br>
Could you explain your use case a bit more and why you are rewrapping the submodules in an <code>nn.Sequential</code> container?</p>

Praveen_K: <p>My use case here is to train <code>Autoencoder</code> and use the model as classifier by adding <code>FC-layer</code> at last.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="111445">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>I’m currently unsure, if you really need to change the model definition or if it would be possible to use the custom model definition.</p>
</blockquote>
</aside>
<p>This was my question. I’m currently unsure how to load the model using custom model definition.</p>
<p>I tried the below code also. But it results in <code>_IncompatibleKeys</code> error as above. Below code is the same code used during custom model training(decoder part not included for obvious reasons).</p>
<pre><code class="lang-auto">class Initialise(nn.Module):
    def __init__(self):
        super(Initialise, self).__init__()
        self.encoder = nn.Sequential(nn.Linear(28*28*3, 256), nn.ReLU())


    def forward(self, x):
        x = self.encoder(x)
        
        return x
		
		
		
model= Initialise()
model.add_module('classifier', nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1)))


</code></pre>
<aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="111445">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>why you are rewrapping the submodules in an <code>nn.Sequential</code> container</p>
</blockquote>
</aside>
<p>I’m using <code>nn.Sequential</code> in <code>encoder</code> and <code>classifier</code> for modules to be in order. and i saw removing the rewrapping wouldn’t solve the problem either <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="111445">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>One fix would be to rename the keys in the <code>state_dict</code>.</p>
</blockquote>
</aside>
<p>I did this and solved the problem. But its not optimum solution, as the code is not generic.</p>
<p>I think the problem comes when i copy the weights of <code>encoder</code>. Here the key name is changed when I copy to train <code>classifier</code> network</p>
<pre><code class="lang-auto">
# copy the weights and bias of encoder
new_classifier = nn.Sequential(*list(net.children())[:-1])

net = new_classifier
# add classifier layer
net.add_module('classifier', nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1)))
</code></pre>
<p>I’m not sure how to copy along with the key name</p>

ptrblck: <p>This model definition:</p>
<pre><code class="lang-python">class Initialise(nn.Module):
    def __init__(self):
        super(Initialise, self).__init__()
        self.encoder = nn.Sequential(nn.Linear(28*28*3, 256), nn.ReLU())


    def forward(self, x):
        x = self.encoder(x)
        
        return x
		
		
		
model= Initialise()
model.add_module('classifier', nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1)))
</code></pre>
<p>won’t use the <code>classifier</code> submodules, as <code>Initialise</code> is not using it in its <code>forward</code> method as can be seen here:</p>
<pre><code class="lang-python">x = torch.randn(1, 2352)
out = model(x)
print(out.shape)
&gt; torch.Size([1, 256])
</code></pre>
<p>You could use a placeholder module and replace it via:</p>
<pre><code class="lang-python">class Initialise(nn.Module):
    def __init__(self):
        super(Initialise, self).__init__()
        self.encoder = nn.Sequential(nn.Linear(28*28*3, 256), nn.ReLU())
        self.classifier = nn.Identity()

    def forward(self, x):
        x = self.encoder(x)
        x = self.classifier(x)
        return x
		
		

model= Initialise()
model.classifier = nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1))

x = torch.randn(1, 2352)
out = model(x)
print(out.shape)
&gt; torch.Size([1, 100])
</code></pre>
<p>You could also save and load the <code>state_dict</code> via:</p>
<pre><code class="lang-python">torch.save(model.state_dict(), 'tmp.pt')
model = Initialise()
model.classifier = nn.Sequential(nn.Linear(256, 100),nn.LogSoftmax(dim=1))
model.load_state_dict(torch.load('tmp.pt'))
</code></pre>
<p>Note however, that the replacement of the <code>nn.Identity()</code> module would still be necessary.<br>
Let me know, if this would work for you.</p>

Praveen_K: <p>Thank you for your reply.</p>
<p>The error persists in the same way. Changing the key names results in matching the keys and loading the model properly.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f9493c02e3194b15df651373efb3306d450d3f1.png" data-download-href="https://discuss.pytorch.org/uploads/default/7f9493c02e3194b15df651373efb3306d450d3f1" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f9493c02e3194b15df651373efb3306d450d3f1_2_690x232.png" alt="image" data-base62-sha1="icCXKWzjvPzaJUIC6h3JtrGHWyB" width="690" height="232" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f9493c02e3194b15df651373efb3306d450d3f1_2_690x232.png, https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f9493c02e3194b15df651373efb3306d450d3f1.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f9493c02e3194b15df651373efb3306d450d3f1.png 2x" data-dominant-color="F1F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">987×332 33.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>As i had mentioned,</p>
<aside class="quote no-group" data-username="Praveen_K" data-post="6" data-topic="111445">
<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/praveen_k/48/27741_2.png" class="avatar"> Praveen_K:</div>
<blockquote>
<p>I think the problem comes when i copy the weights of <code>encoder</code>. Here the key name is changed when I copy to train <code>classifier</code> network</p>
</blockquote>
</aside>
<p>While copying the <code>encoder</code> for further training as <code>classifier</code>, I’m not copying the key names and this is resulting in error. The name before copying was <code>encoder</code> after it will be <code>0</code>. Hence <code>0.0.weight</code>, <code>0.0.bias</code></p>
<p>P.S : Sorry for pasting an image instead of pasting the code <img src="https://discuss.pytorch.org/images/emoji/apple/no_mouth.png?v=12" title=":no_mouth:" class="emoji" alt=":no_mouth:" loading="lazy" width="20" height="20"></p>