Error when saving and loading GAN models

username: abedshantti
cooked: <p>I am attempting to save my trained GAN models so that I can either resume training later, or use my saved models for inference.</p>
<p>My <strong>first</strong> attempt for saving the model was by saving the entire model as:</p>
<pre><code class="lang-auto">model = myGAN(epochs = 10, additional_args)
model.fit(data)

torch.save(model, "model.pth")
</code></pre>
<p>But this prompts the following error:</p>
<pre><code class="lang-auto">**PicklingError** : Can't pickle &lt;class 'Model.myGAN.columnInfo'&gt;: attribute lookup columnInfo on Model.myGAN failed
</code></pre>
<p><code>columnInfo</code> is a named tuple that I am using in my model which has access to the my data’s feature descriptions, but apparently it cannot be saved.</p>
<p>My <strong>second</strong> attempt for saving and loading the model was using the recommended Pytorch way:</p>
<pre><code class="lang-auto">model = myGAN(epochs = 10, additional_args)
model.fit(data)

torch.save(model.state_dict(), "model.pth")

model2 = myGAN(epochs = 5, additional_args)
model2.load_state_dict(torch.load("model.pth"))
</code></pre>
<p>But that returns the following error:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for myGAN:
	Unexpected key(s) in state_dict: 
</code></pre>
<p>I believe this happens because some of the unexpected keys are only initialised when the model begins training, and thereby are not recognised by model2 before that. Are there any workarounds for this errors?</p>
<p>Edit1: Loading the model parameters only works if it is done on the same model I initialised and trained. e.g.:</p>
<pre><code class="lang-auto">model = myGAN(epochs = 10, additional_args)
model.fit(data)

torch.save(model.state_dict(), "model.pth")

model.load_state_dict(torch.load("model.pth"))
</code></pre>

username: InnovArul
cooked: <aside class="quote no-group" data-username="abedshantti" data-post="1" data-topic="137122">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/8797f3/48.png" class="avatar"> abedshantti:</div>
<blockquote>
<p><code>load_state_dict</code></p>
</blockquote>
</aside>
<p>You can try setting <code>strict=False</code> parameter while using <code>model.load_state_dict()</code> function.<br>
That might help!</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict">load_state_dict documentation</a></p>

username: abedshantti
cooked: <p>Unfortunately using <code>strict = False</code> did not work for me. However, I found a solution, but the workaround is a bit awkward:</p>
<pre><code class="lang-auto">model = myGAN(epochs = 10, additional_args)
model.fit(data)

torch.save({"epoch" : 10,
            "model_state_dict" : model.state_dict(),
            "G_optimizer_state_dict" : self.G_optimizer.state_dict(),
            "D_optimizer_state_dict" : self.D_optimizer.state_dict() }, 
            "model.pth")

model2 = myGAN()
model2 = myGAN.fit(data, epochs = 0)

checkpoint = torch.load("model.pth")
model2.load_state_dict(checkpoint['model_state_dict'])
model2.G_optimizer.load_state_dict(checkpoint['G_optimizer_state_dict'])
model2.D_optimizer.load_state_dict(checkpoint['D_optimizer_state_dict'])
 
model2.fit(data, epochs = 10)
</code></pre>
<p>I moved the epochs inside the fit function, but I still have to call fit with 0 epochs before loading the model to accept the saved parameters.</p>

username: InnovArul
cooked: <p>I am not sure why it didn’t work for you.<br>
But the below code works.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.m1 = nn.Linear(5,6)
        self.m2 = nn.Linear(6,7)

    def create_dynamic_linear(self):
        self.m3 = nn.Linear(3,3)

    def forward(self,):
        pass

if __name__ == '__main__':
    model1 = model()
    model1.create_dynamic_linear()
    torch.save(model1.state_dict(), 'model.pth')

    model2 = model()
    miss, unexp = model2.load_state_dict(torch.load('model.pth'), strict=False)

    print(miss, unexp)


</code></pre>