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

.. only:: html

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

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

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

.. _sphx_glr_auto_examples_drop_in_replacement_tutorial.py:


======================================================
From ResNets to Momentum ResNets 1)
======================================================

This is a tutorial to use the transform_to_momentumnet
method:

Michael E. Sander, Pierre Ablin, Mathieu Blondel,
Gabriel Peyre. Momentum Residual Neural Networks.
Proceedings of the 38th International Conference 
on Machine Learning, PMLR 139:9276-9287

.. GENERATED FROM PYTHON SOURCE LINES 16-22

.. code-block:: default


    # Authors: Michael Sander, Pierre Ablin
    # License: MIT
    from torch import nn
    from momentumnet import transform_to_momentumnet


.. GENERATED FROM PYTHON SOURCE LINES 23-25

Let us define a toy Neural Network
###################################

.. GENERATED FROM PYTHON SOURCE LINES 25-77

.. code-block:: default



    class ResBlock(nn.Module):
        def __init__(self, functions):
            super(ResBlock, self).__init__()
            self.functions = functions

        def forward(self, x):

            for f in self.functions:
                x = x + f(x)

            return x


    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.res_layer1 = ResBlock(
                nn.Sequential(
                    *[
                        nn.Sequential(
                            nn.Linear(2, 10), nn.Tanh(), nn.Linear(10, 2)
                        )
                        for _ in range(3)
                    ]
                )
            )
            self.l1 = nn.Linear(2, 4)
            self.layer2 = nn.Sequential(
                *[
                    nn.Sequential(nn.Linear(4, 100), nn.ReLU(), nn.Linear(100, 4))
                    for _ in range(4)
                ]
            )
            self.l2 = nn.Linear(4, 8)

            self.fc = nn.Linear(8, 10)

        def forward(self, x):

            out = self.res_layer1(x)  # Residual
            out = self.l1(out)
            out = self.layer2(out)  # Not Residual but same dimensions
            out = self.l2(out)
            out = self.fc(out)

            return out


    net = Net()


.. GENERATED FROM PYTHON SOURCE LINES 78-80

We want to transform it into its Momentum version
##################################################

.. GENERATED FROM PYTHON SOURCE LINES 82-86

The first layer 'res_layer1' preserves dimension and is residual.
It can be accessed through net.res_layer_1.functions so we will specify
this attribute as the "sub_layers" parameter.
One can transform this residual block into a momentum one as follow:

.. GENERATED FROM PYTHON SOURCE LINES 86-96

.. code-block:: default


    mnet1 = transform_to_momentumnet(
        net,
        ["res_layer1.functions"],  # attributes of the sublayers in net
        gamma=0.9,
        use_backprop=False,
        is_residual=True,
        keep_first_layer=False,
    )


.. GENERATED FROM PYTHON SOURCE LINES 97-101

Note that layer2 is not residual but also preserves dimensions.
It can be accessed through net.layer_2 so we will specify
this attribute as the "sub_layers" parameter.
One can transform it in the same way setting is_residual to False.

.. GENERATED FROM PYTHON SOURCE LINES 101-111

.. code-block:: default


    mnet = transform_to_momentumnet(
        mnet1,
        ["layer2"],
        gamma=0.9,
        use_backprop=False,
        is_residual=False,
        keep_first_layer=False,
    )


.. GENERATED FROM PYTHON SOURCE LINES 112-113

net, mnet1, and mnet have the same parameters.


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

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


.. _sphx_glr_download_auto_examples_drop_in_replacement_tutorial.py:


.. only :: html

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



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

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



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

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


.. only:: html

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

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