.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "beginner/basics/saveloadrun_tutorial.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_beginner_basics_saveloadrun_tutorial.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_beginner_basics_saveloadrun_tutorial.py:


`Learn the Basics <intro.html>`_ ||
`Quickstart <quickstart_tutorial.html>`_ ||
`Tensors <tensorqs_tutorial.html>`_ ||
`Datasets & DataLoaders <data_tutorial.html>`_ ||
`Transforms <transforms_tutorial.html>`_ ||
`Build Model <buildmodel_tutorial.html>`_ ||
`Autograd <autogradqs_tutorial.html>`_ ||
`Optimization <optimization_tutorial.html>`_ ||
**Save & Load Model**

Save and Load the Model
============================

In this section we will look at how to persist model state with saving, loading and running model predictions.

.. GENERATED FROM PYTHON SOURCE LINES 17-22

.. code-block:: default


    import torch
    import torchvision.models as models









.. GENERATED FROM PYTHON SOURCE LINES 23-28

Saving and Loading Model Weights
--------------------------------
PyTorch models store the learned parameters in an internal
state dictionary, called ``state_dict``. These can be persisted via the ``torch.save``
method:

.. GENERATED FROM PYTHON SOURCE LINES 28-32

.. code-block:: default


    model = models.vgg16(weights='IMAGENET1K_V1')
    torch.save(model.state_dict(), 'model_weights.pth')





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/vgg16-397923af.pth

      0%|          | 0.00/528M [00:00<?, ?B/s]
      3%|3         | 16.2M/528M [00:00<00:03, 169MB/s]
      6%|6         | 33.6M/528M [00:00<00:02, 176MB/s]
     10%|9         | 51.4M/528M [00:00<00:02, 181MB/s]
     13%|#3        | 69.0M/528M [00:00<00:02, 182MB/s]
     16%|#6        | 86.5M/528M [00:00<00:02, 182MB/s]
     20%|#9        | 104M/528M [00:00<00:02, 183MB/s] 
     23%|##3       | 122M/528M [00:00<00:02, 182MB/s]
     26%|##6       | 139M/528M [00:00<00:02, 183MB/s]
     30%|##9       | 157M/528M [00:00<00:02, 183MB/s]
     33%|###3      | 174M/528M [00:01<00:02, 184MB/s]
     36%|###6      | 192M/528M [00:01<00:01, 184MB/s]
     40%|###9      | 210M/528M [00:01<00:01, 183MB/s]
     43%|####3     | 227M/528M [00:01<00:01, 180MB/s]
     46%|####6     | 244M/528M [00:01<00:01, 179MB/s]
     50%|####9     | 262M/528M [00:01<00:01, 179MB/s]
     53%|#####2    | 279M/528M [00:01<00:01, 179MB/s]
     56%|#####6    | 296M/528M [00:01<00:01, 179MB/s]
     59%|#####9    | 313M/528M [00:01<00:01, 178MB/s]
     63%|######2   | 330M/528M [00:01<00:01, 178MB/s]
     66%|######5   | 347M/528M [00:02<00:01, 178MB/s]
     69%|######9   | 364M/528M [00:02<00:00, 178MB/s]
     72%|#######2  | 382M/528M [00:02<00:00, 178MB/s]
     76%|#######5  | 399M/528M [00:02<00:00, 179MB/s]
     79%|#######8  | 416M/528M [00:02<00:00, 179MB/s]
     82%|########2 | 433M/528M [00:02<00:00, 179MB/s]
     85%|########5 | 450M/528M [00:02<00:00, 179MB/s]
     89%|########8 | 467M/528M [00:02<00:00, 179MB/s]
     92%|#########1| 484M/528M [00:02<00:00, 179MB/s]
     95%|#########5| 502M/528M [00:02<00:00, 179MB/s]
     98%|#########8| 519M/528M [00:03<00:00, 179MB/s]
    100%|##########| 528M/528M [00:03<00:00, 180MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 33-35

To load model weights, you need to create an instance of the same model first, and then load the parameters
using ``load_state_dict()`` method.

.. GENERATED FROM PYTHON SOURCE LINES 35-40

.. code-block:: default


    model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
    model.load_state_dict(torch.load('model_weights.pth'))
    model.eval()





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    VGG(
      (features): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (6): ReLU(inplace=True)
        (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (8): ReLU(inplace=True)
        (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (11): ReLU(inplace=True)
        (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): ReLU(inplace=True)
        (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (15): ReLU(inplace=True)
        (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (18): ReLU(inplace=True)
        (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (20): ReLU(inplace=True)
        (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (22): ReLU(inplace=True)
        (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (25): ReLU(inplace=True)
        (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (27): ReLU(inplace=True)
        (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (29): ReLU(inplace=True)
        (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      )
      (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
      (classifier): Sequential(
        (0): Linear(in_features=25088, out_features=4096, bias=True)
        (1): ReLU(inplace=True)
        (2): Dropout(p=0.5, inplace=False)
        (3): Linear(in_features=4096, out_features=4096, bias=True)
        (4): ReLU(inplace=True)
        (5): Dropout(p=0.5, inplace=False)
        (6): Linear(in_features=4096, out_features=1000, bias=True)
      )
    )



.. GENERATED FROM PYTHON SOURCE LINES 41-42

.. note:: be sure to call ``model.eval()`` method before inferencing to set the dropout and batch normalization layers to evaluation mode. Failing to do this will yield inconsistent inference results.

.. GENERATED FROM PYTHON SOURCE LINES 44-49

Saving and Loading Models with Shapes
-------------------------------------
When loading model weights, we needed to instantiate the model class first, because the class
defines the structure of a network. We might want to save the structure of this class together with
the model, in which case we can pass ``model`` (and not ``model.state_dict()``) to the saving function:

.. GENERATED FROM PYTHON SOURCE LINES 49-52

.. code-block:: default


    torch.save(model, 'model.pth')








.. GENERATED FROM PYTHON SOURCE LINES 53-54

We can then load the model like this:

.. GENERATED FROM PYTHON SOURCE LINES 54-57

.. code-block:: default


    model = torch.load('model.pth')








.. GENERATED FROM PYTHON SOURCE LINES 58-59

.. note:: This approach uses Python `pickle <https://docs.python.org/3/library/pickle.html>`_ module when serializing the model, thus it relies on the actual class definition to be available when loading the model.

.. GENERATED FROM PYTHON SOURCE LINES 61-65

Related Tutorials
-----------------
- `Saving and Loading a General Checkpoint in PyTorch <https://pytorch.org/tutorials/recipes/recipes/saving_and_loading_a_general_checkpoint.html>`_
- `Tips for loading an nn.Module from a checkpoint <https://pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html?highlight=loading%20nn%20module%20from%20checkpoint>`_


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  8.725 seconds)


.. _sphx_glr_download_beginner_basics_saveloadrun_tutorial.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example


    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: saveloadrun_tutorial.py <saveloadrun_tutorial.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: saveloadrun_tutorial.ipynb <saveloadrun_tutorial.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_