Comparison between saving the whole model, saving only state_dict and torchscript

LewsTherin: <p>I’m trying to figure out what’s the best way to save a model trained with Pytorch and load it for inference, and I was wondering about the different possible approaches.</p>
<p>Let’s say I successfully train a model, as far as I understand I can use:</p>
<ul>
<li><em><strong>Complete Model Saving:</strong></em></li>
</ul>
<pre><code class="lang-auto"># save the model
torch.save(model, saved_model_path)
# load model directly with
loaded_model = torch.load(saved_model_path)
# use it for inference
output = loaded_model(input)
</code></pre>
<ul>
<li><em><strong>State Dict Saving:</strong></em></li>
</ul>
<pre><code class="lang-auto"># save only the state_dict after training
torch.save(model.state_dict(), saved_model_path)
# need to create an instance of the model with the same architecture and then load the parameters using
model = SomeModelConstructor()
model.load_state_dict(state_dict)
# use it for inference
output = loaded_model(input)
</code></pre>
<p>And I seem to understand saving only the state_dict has the advantage of reducing the file size, while the con is that I have to recreate the model instance. Is that correct?</p>
<p>Also, I see I can also export the model to torchscript and load it:</p>
<pre><code class="lang-auto"># export to torchscript and save locally
model_scripted = torch.jit.script(model)
torch.jit.save(model_scripted, scripted_model_path)

# load the scripted model
loaded_scripted_model = torch.jit.load(scripted_model_path)

# use for evaluation
loaded_scripted_model.eval()
with torch.no_grad():    
    output = loaded_scripted_model(input)
</code></pre>
<p>Assuming I’m exporting the model to script from Python and loading it back always from Python, are there any advantages in using the scripted model instead of the saved one (for example, inference speed because the scripted version is better optimized or something similar?)</p>

ptrblck: <aside class="quote no-group" data-username="LewsTherin" data-post="1" data-topic="196033">
<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/l/5f9b8f/48.png" class="avatar"> LewsTherin:</div>
<blockquote>
<p>saving only the state_dict has the advantage of reducing the file size, while the con is that I have to recreate the model instance. Is that correct?</p>
</blockquote>
</aside>
<p>Yes, you would need to recreate the model instance when saving the <code>state_dict</code> only, but note that saving the entire model requires you to exactly recreate the same file structure (including the source code of the model) on the target system, which is quite brittle and you will this see a lot of topics in this discussion board of users asking how to load the entire model now after moving files around or changing function signatures as it fails.</p>
<p>Saving and loading the <code>state_dict</code> is thus the recommended approach.</p>
<aside class="quote no-group" data-username="LewsTherin" data-post="1" data-topic="196033">
<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/l/5f9b8f/48.png" class="avatar"> LewsTherin:</div>
<blockquote>
<p>are there any advantages in using the scripted model instead of the saved one</p>
</blockquote>
</aside>
<p>Scripted models can be optimized by <code>TorchScript</code> and loaded into the C++ <code>libtorch</code> runtime directly. However, <code>TorchScript</code> is now in maintenance mode and I would not depend on it anymore.</p>

LewsTherin: <p>Thank you very much for your reply! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Just a couple of other question:</p>
<blockquote>
<p>Scripted models can be optimized by <code>TorchScript</code> and loaded into the C++ <code>libtorch</code> runtime directly</p>
</blockquote>
<p>So, as I was mentioning, if I simply load back the scripted model from Python and use it for inference, there is no such optimization?</p>
<blockquote>
<p>However, <code>TorchScript</code> is now in maintenance mode and I would not depend on it anymore.</p>
</blockquote>
<p>Could you please elaborate on this? Torchscript will be discontinued?</p>

ptrblck: <aside class="quote no-group" data-username="LewsTherin" data-post="3" data-topic="196033">
<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/l/5f9b8f/48.png" class="avatar"> LewsTherin:</div>
<blockquote>
<p>Could you please elaborate on this? Torchscript will be discontinued?</p>
</blockquote>
</aside>
<p>Yes, that’s my understanding as <code>TorchDynamo</code> will be developed instead.</p>
<aside class="quote no-group" data-username="LewsTherin" data-post="3" data-topic="196033">
<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/l/5f9b8f/48.png" class="avatar"> LewsTherin:</div>
<blockquote>
<p>if I simply load back the scripted model from Python and use it for inference, there is no such optimization?</p>
</blockquote>
</aside>
<p><code>TorchScript</code> would also optimize the model in your Python script, but I would generally not recommend relying on it anymore. Use <code>torch.compile</code> or the beta <code>torch.export</code> support and try to run your model.</p>