About Model Saving

xiangyu_tang: <p>During the model training, I saved the parameters of the model in the following way:</p>
<pre><code class="lang-auto">torch.save(net.state_dict(),path)
</code></pre>
<p>After the training, I wanted to get a complete model, so I ran the following Python script:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.models as models


class MobileNet(nn.Module):
    def __init__(self, num_classes, mode, pretrained=True):
        super(MobileNet, self).__init__()
        assert mode in ["mobilenet_v2", "mobilenet_v3_small", "mobilenet_v3_large"]
        self.num_classes = num_classes
        self.mode = mode

        if mode == "mobilenet_v2":
            model = models.mobilenet_v2(pretrained=pretrained)
            model.classifier[1] = nn.Linear(in_features=1280, out_features=self.num_classes, bias=True)

        elif mode == "mobilenet_v3_small":
            model = models.mobilenet_v3_small(pretrained=pretrained)
            model.classifier[3] = nn.Linear(in_features=1024, out_features=self.num_classes, bias=True)

        elif mode == "mobilenet_v3_large":
            model = models.mobilenet_v3_large(pretrained=pretrained)
            model.classifier[3] = nn.Linear(in_features=1280, out_features=self.num_classes, bias=True)

        self.model = model

    def forward(self, x):
        x = self.model(x)
        return x


if __name__ == "__main__":
    our_image = torch.rand(size=(2, 3, 224, 224))
    our_model = MobileNet(num_classes=4, mode="mobilenet_v2")
    our_model = torch.nn.DataParallel(our_model)
    our_model.load_state_dict(torch.load(path))
    torch.save(our_model, "./our_model.pt")

</code></pre>
<p>However, when I used the complete model I had saved, I got it wrong. The Python script used and the error is as follows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/3/8390543a706de8be61a8d4c7224fcbb319a2b642.png" data-download-href="https://discuss.pytorch.org/uploads/default/8390543a706de8be61a8d4c7224fcbb319a2b642" title="1639019590(1)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/8390543a706de8be61a8d4c7224fcbb319a2b642_2_690x356.png" alt="1639019590(1)" data-base62-sha1="iLRLS23pQ6VPG4dNJ1LTcZbMrQK" width="690" height="356" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/8390543a706de8be61a8d4c7224fcbb319a2b642_2_690x356.png, https://discuss.pytorch.org/uploads/default/original/3X/8/3/8390543a706de8be61a8d4c7224fcbb319a2b642.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/3/8390543a706de8be61a8d4c7224fcbb319a2b642.png 2x" data-dominant-color="323133"><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">1639019590(1)</span><span class="informations">967×500 47.6 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><br>
I was rather puzzled by the error. There is an error indicating that I need to provide the script for building the model, but I have saved the complete model, so I don’t think I need the script for building the model.</p>

ptrblck: <p>Based on the error message it seems as if you’ve saved the model object instead of it’s <code>state_dict</code>.<br>
Are you sure you’ve used <code>torch.save(net.state_dict(),path)</code> and not <code>torch.save(net,path)</code>?<br>
In the latter case, you would need to restore the file structure again to be able to load the model object.</p>

xiangyu_tang: <p>Thank you for your reply. I’m sure I’m using <code>torch.save (net,path)</code>. The Python script I use to save the full model is as follows:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9d7e19ef478935c23e2940b4d58f3d22c055db2.png" alt="1639027375(1)" data-base62-sha1="zEdpEv2nlvhznI75bqzWBtFzrpw" width="516" height="113"></p>

ptrblck: <p>Ah, OK this would explain the error, as this method relies on the same file structure to be able to load the model again.<br>
Check these <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model">docs</a> which explain the disadvantage of saving the model object instead of its <code>state_dict</code>:</p>
<blockquote>
<p>This save/load process uses the most intuitive syntax and involves the least amount of code. Saving a model in this way will save the entire module using Python’s <a href="https://docs.python.org/3/library/pickle.html">pickle</a> module. The disadvantage of this approach is that the serialized data is bound to the specific classes and the exact directory structure used when the model is saved. The reason for this is because pickle does not save the model class itself. Rather, it saves a path to the file containing the class, which is used during load time. Because of this, your code can break in various ways when used in other projects or after refactors.</p>
</blockquote>

xiangyu_tang: <p>Thank you very much. I think I get it. It just saves the path. So what do I do if I want to save a complete model and then use it in other projects?</p>

ptrblck: <p>I think you would need to copy/paste the entire source code with the stored model, but I have to admit not to use this approach as it can easily break and stick to saving the <code>state_dict</code>s instead.</p>

xiangyu_tang: <p>Ok, I see. Thank you very much.</p>

xiangyu_tang: <p>I want to ask you one more question. So if I save the complete model including the structure of the model directly when I train the model, I should be able to save the structure of the model to apply to other projects, instead of just storing parameters and the path of the structure of the model. Or at what point does it save only parameters and model paths?</p>

ptrblck: <p>The common two approaches are mentioned in the linked doc and are basically:</p>
<ul>
<li>store the <code>state_dict</code>s and keep the model source code around. You can refactor the code and as long as the model’s <code>__init__</code> method creates the same submodules you will be able to load the <code>state_dict</code> into the model. You could change the <code>forward</code> method of the model as the <code>state_dict</code> doesn’t have any knowledge about the forward pass.</li>
<li>store the model object directly and make sure the project structure is equal in your other setup where you would like to load the model</li>
</ul>
<p>You could also check the new <a href="https://pytorch.org/docs/stable/package.html"><code>torch.package</code></a> module, which might be useful for your use case.</p>

xiangyu_tang: <p>Ok, I see. Thank you for your help.</p>