Loading Weights From saved model

username: nbansal90
cooked: <p>Hello,</p>
<p>In my case I have a saved model named <code>best_model.pkl</code>, which inherently contains multiple models belonging to <code>pose</code>, <code>depth</code> and other models. I was looking to use the existing <code>best_model.pkl</code> to extract the individual model and save them as separate <code>.pth</code> files.</p>
<p>I am doing something on this line:</p>
<pre><code class="lang-auto">model_to_save = torch.load('best_model.pkl')
print (model_to_save.keys())
</code></pre>
<p>which gives the following output:</p>
<pre><code class="lang-auto">dict_keys(['epoch', 'model_state', 'optimizer_state', 'scheduler_state', 'best_iou'])
</code></pre>
<p>Of which I select the <code>model_state</code>, which has each of the individual models. Now I am looking to save each of the individual models by the following code.</p>
<pre><code class="lang-auto">models = ["depth", "pose_encoder", "pose", "encoder"]
    for model_name in models:
        save_path = os.path.join("{}.pth".format(model_name))
        to_save = model_to_save['model_state']['models'][model_name].state_dict()
        torch.save(to_save, save_path)
</code></pre>
<p>But I receive following error:</p>
<pre><code class="lang-auto">  File "save.py", line 10, in save_monodepth_models
    to_save = model_to_save['model_state']['models'][model_name].state_dict()
KeyError: 'models'
</code></pre>
<p>I see that <code>type(model_to_save['model_state'])</code> is an <code>&lt;class 'collections.OrderedDict'&gt;</code> and keys inside the <code>model_to_save['model_state']</code> is:</p>
<pre><code class="lang-auto">odict_keys(['models.encoder.encoder.conv1.weight', 'models.encoder.encoder.bn1.weight', 'models.encoder.encoder.bn1.bias', 'models.encoder.encoder.bn1.running_mean', 'models.enco
der.encoder.bn1.running_var', 'models.encoder.encoder.bn1.num_batches_tracked', 'models.encoder.encoder.layer1.0.conv1.weight', 'models.encoder.encoder.layer1.0.bn1.weight', 'mod
els.encoder.encoder.layer1.0.bn1.bias', 'models.encoder.encoder.layer1.0.bn1.running_mean', 'models.encoder.encoder.layer1.0.bn1.running_var', 'models.encoder.encoder.layer1.0.bn
1.num_batches_tracked', 'models.encoder.encoder.layer1.0.conv2.weight', 'models.encoder.encoder.layer1.0.bn2.weight', 'models.encoder.encoder.layer1.0.bn2.bias', 'models.encoder.
encoder.layer1.0.bn2.running_mean', 'models.encoder.encoder.layer1.0.bn2.running_var', 'models.encoder.encoder.layer1.0.bn2.num_batches_tracked', 'models.encoder.encoder.layer1.0
.conv3.weight', 'models.encoder.encoder.layer1.0.bn3.weight', 'models.encoder.encoder.layer1.0.bn3.bias', 'models.encoder.encoder.layer1.0.bn3.running_mean', 'models.encoder.enco
der.layer1.0.bn3.running_var',................
</code></pre>
<p>Not exactly sure how should I handle this case. Any suggestion would be helpful!</p>
<p>Regards!</p>

username: ptrblck
cooked: <p>It seems that <code>model_to_save['model_state']</code> contains a <code>state_dict</code>, which is the recommended way to serialize the models.<br>
To restore the model, create a model instance first and load the corresponding <code>state_dict</code> afterwards.<br>
I would not recommend to save the model directly via <code>torch.save(model)</code> but to stick to saving the <code>state_dict</code> via <code>torch.save(model.state_dict())</code> (or in your case you could also wrap it into more <code>dict</code>s if needed) as the former might easily break (e.g. you would need to restore the same file structure etc.).</p>