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

.. only:: html

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

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

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

.. _sphx_glr_recipes_recipes_warmstarting_model_using_parameters_from_a_different_model.py:


Warmstarting model using parameters from a different model in PyTorch
=====================================================================
Partially loading a model or loading a partial model are common
scenarios when transfer learning or training a new complex model.
Leveraging trained parameters, even if only a few are usable, will help
to warmstart the training process and hopefully help your model converge
much faster than training from scratch.

Introduction
------------
Whether you are loading from a partial ``state_dict``, which is missing
some keys, or loading a ``state_dict`` with more keys than the model
that you are loading into, you can set the strict argument to ``False``
in the ``load_state_dict()`` function to ignore non-matching keys.
In this recipe, we will experiment with warmstarting a model using
parameters of a different model.

Setup
-----
Before we begin, we need to install ``torch`` if it isn’t already
available.

.. code-block:: sh

   pip install torch
   

.. GENERATED FROM PYTHON SOURCE LINES 33-47

Steps
-----

1. Import all necessary libraries for loading our data
2. Define and initialize the neural network A and B
3. Save model A
4. Load into model B

1. Import necessary libraries for loading our data
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For this recipe, we will use ``torch`` and its subsidiaries ``torch.nn``
and ``torch.optim``.


.. GENERATED FROM PYTHON SOURCE LINES 47-53

.. code-block:: default


    import torch
    import torch.nn as nn
    import torch.optim as optim



.. GENERATED FROM PYTHON SOURCE LINES 54-62

2. Define and initialize the neural network A and B
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For sake of example, we will create a neural network for training
images. To learn more see the Defining a Neural Network recipe. We will
create two neural networks for sake of loading one parameter of type A
into type B.


.. GENERATED FROM PYTHON SOURCE LINES 62-106

.. code-block:: default


    class NetA(nn.Module):
        def __init__(self):
            super(NetA, self).__init__()
            self.conv1 = nn.Conv2d(3, 6, 5)
            self.pool = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(6, 16, 5)
            self.fc1 = nn.Linear(16 * 5 * 5, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 10)

        def forward(self, x):
            x = self.pool(F.relu(self.conv1(x)))
            x = self.pool(F.relu(self.conv2(x)))
            x = x.view(-1, 16 * 5 * 5)
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x

    netA = NetA()

    class NetB(nn.Module):
        def __init__(self):
            super(NetB, self).__init__()
            self.conv1 = nn.Conv2d(3, 6, 5)
            self.pool = nn.MaxPool2d(2, 2)
            self.conv2 = nn.Conv2d(6, 16, 5)
            self.fc1 = nn.Linear(16 * 5 * 5, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 10)

        def forward(self, x):
            x = self.pool(F.relu(self.conv1(x)))
            x = self.pool(F.relu(self.conv2(x)))
            x = x.view(-1, 16 * 5 * 5)
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x

    netB = NetB()



.. GENERATED FROM PYTHON SOURCE LINES 107-110

3. Save model A
~~~~~~~~~~~~~~~~~~~


.. GENERATED FROM PYTHON SOURCE LINES 110-117

.. code-block:: default


    # Specify a path to save to
    PATH = "model.pt"

    torch.save(netA.state_dict(), PATH)



.. GENERATED FROM PYTHON SOURCE LINES 118-126

4. Load into model B
~~~~~~~~~~~~~~~~~~~~~~~~

If you want to load parameters from one layer to another, but some keys
do not match, simply change the name of the parameter keys in the
state_dict that you are loading to match the keys in the model that you
are loading into.


.. GENERATED FROM PYTHON SOURCE LINES 126-130

.. code-block:: default


    netB.load_state_dict(torch.load(PATH), strict=False)



.. GENERATED FROM PYTHON SOURCE LINES 131-143

You can see that all keys matched successfully!

Congratulations! You have successfully warmstarted a model using
parameters from a different model in PyTorch.

Learn More
----------

Take a look at these other recipes to continue your learning:

- `Saving and loading multiple models in one file using PyTorch <https://pytorch.org/tutorials/recipes/recipes/saving_multiple_models_in_one_file.html>`__
- `Saving and loading models across devices in PyTorch <https://pytorch.org/tutorials/recipes/recipes/save_load_across_devices.html>`__


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

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


.. _sphx_glr_download_recipes_recipes_warmstarting_model_using_parameters_from_a_different_model.py:

.. only:: html

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


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

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

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

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


.. only:: html

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

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