Saving model AND optimiser AND scheduler

ONTDave: <p>Hi,<br>
I want to able to have a model/optimiser/scheduler object - which I can hot plug and play.<br>
So for example, have a list of such objects, load to gpu in turn, do some training, switch objects.<br>
Maybe then load some earlier ones and pick up training where we left off last time.<br>
I’d like to be able to easily (deep) copy these objects, and save/load to disk.<br>
Note - some models or optimisers or schedulers may be different in these different objects.</p>
<p>One idea - use the torch.save(model) - this will pickle the model class and reproduce the object and load the state_dict, but will it restore the scheduler/optimiser?<br>
Can I use torch.save with these as well to reproduce them later?<br>
How do I control which model they are attached to?</p>
<p>And how can I duplicate the whole model/optimizer/scheduler whilst in memory?</p>
<p>Any suggestions on the best way to do this please?<br>
Many thanks in advance!</p>

dpernes: <p>You can create a dictionary with everything you need and save it using <code>torch.save()</code>. Example:</p>
<pre><code class="lang-auto">checkpoint = { 
    'epoch': epoch,
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'lr_sched': lr_sched}
torch.save(checkpoint, 'checkpoint.pth')
</code></pre>
<p>Then you can load the checkpoint doing <code>checkpoint = torch.load('checkpoint.pth')</code><br>
More info here: <a href="https://discuss.pytorch.org/t/loading-a-saved-model-for-continue-training/17244" class="inline-onebox">Loading a saved model for continue training</a></p>

ONTDave: <p>Hi, many thanks for your quick reply! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
So I’ve seen that. But the problem with that is it only copies the models state dict.<br>
If different objects say have different types of models, how do I know which model to create before giving it the saved state dict?<br>
Same for optimizers and schedulers.</p>

dpernes: <p>Then you can save the <code>model</code> itself, i.e. without calling <code>state_dict()</code>. This will dump the whole content of the variable <code>model</code> into a pickle file, which leads to a larger file than in the previous case. The same applies to the <code>optimizer</code>. The scheduler does not have a <code>state_dict()</code>, so the whole variable is saved anyways.</p>

ONTDave: <p>So I save the model. And then the optimizer. They will be in two separate files? Will they be connected to each other correctly on reloading. And if I want to say create 10 duplicates of this “group” is there a way to do that?<br>
Many thanks <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

dpernes: <blockquote>
<p>They will be in two separate files?</p>
</blockquote>
<p>You can save them in separate files or wrap them in a dictionary like the one I showed (where you remove the calls to <code>state_dict()</code>).</p>
<blockquote>
<p>Will they be connected to each other correctly on reloading.</p>
</blockquote>
<p>There is no loss of information when you save the whole variable instead of saving the <code>state_dict()</code> only, so I assume that the <code>model</code> and the <code>optimizer</code> will remain “connected” after you reload them (although I’ve never tried it).</p>
<blockquote>
<p>And if I want to say create 10 duplicates of this “group” is there a way to do that?</p>
</blockquote>
<p>I’m not sure if I understand what you mean with a “duplicate”… Do you mean, for instance, saving a checkpoint of the training after every epoch? You may save each checkpoint file with a different name and that’s it.</p>

ONTDave: <p>If I store them in the same file - it only stores the state_dicts for each. It will not pickle the object. Problems then if different “objects” have different models.<br>
Separate files sounds troublesome. Which order do I load them in and do they still work for training.<br>
Imagine one of a group of 10 model/optimizer/scheduler does particularly well after round one - where each has had an hour on the gpu.  This may be because I gave it a favourable set of hyper parameters.<br>
I want to take this group, duplicate it x10, and then run each of those with slightly different hyperparameters.<br>
In essence - I want to be able to duplicate the whole model/optimiser/scheduler state.</p>

dpernes: <blockquote>
<p>If I store them in the same file - it only stores the state_dicts for each. It will not pickle the object.</p>
</blockquote>
<p>This is not true. As long as you do not call <code>state_dict()</code>, it will save the whole variable. Please try the following.</p>
<p>For saving:</p>
<pre><code class="lang-auto">checkpoint = { 
    'epoch': epoch,
    'model': model,
    'optimizer': optimizer,
    'lr_sched': lr_sched}
torch.save(checkpoint, 'checkpoint.pth')
</code></pre>
<p>For loading:</p>
<pre><code class="lang-auto">checkpoint = torch.load('checkpoint.pth')
epoch = checkpoint['epoch']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
lr_sched = checkpoint['lr_sched']
</code></pre>
<p>What’s the problem with this approach?</p>

ONTDave: <p>Oh I see what you mean. I’m sorry - me stupid! <img src="https://discuss.pytorch.org/images/emoji/apple/crazy_face.png?v=9" title=":crazy_face:" class="emoji" alt=":crazy_face:"><br>
Your help is much appreciated! <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"></p>

ONTDave: <p>Oh, sorry - one last thing - wondered about copying the whole thing say x10?<br>
Do I just reload the same thing from disk multiple times?</p>

baromri: <p>I’ve been working on a checkpoint helper for such use case. It is still WIP but check it out: <a href="https://pypi.org/project/pytorchcheckpoint/" rel="nofollow noopener">https://pypi.org/project/pytorchcheckpoint/</a></p>

dpernes: <p>You can reload the same pickle file as many times as you need.</p>