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

.. only:: html

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

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

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

.. _sphx_glr_beginner_basics_optimization_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** ||
`Save & Load Model <saveloadrun_tutorial.html>`_

Optimizing Model Parameters
===========================

Now that we have a model and data it's time to train, validate and test our model by optimizing its parameters on
our data. Training a model is an iterative process; in each iteration the model makes a guess about the output, calculates
the error in its guess (*loss*), collects the derivatives of the error with respect to its parameters (as we saw in
the `previous section  <autograd_tutorial.html>`_), and **optimizes** these parameters using gradient descent. For a more
detailed walkthrough of this process, check out this video on `backpropagation from 3Blue1Brown <https://www.youtube.com/watch?v=tIeHLnjs5U8>`__.

Prerequisite Code
-----------------
We load the code from the previous sections on `Datasets & DataLoaders <data_tutorial.html>`_
and `Build Model  <buildmodel_tutorial.html>`_.

.. GENERATED FROM PYTHON SOURCE LINES 26-70

.. code-block:: default


    import torch
    from torch import nn
    from torch.utils.data import DataLoader
    from torchvision import datasets
    from torchvision.transforms import ToTensor

    training_data = datasets.FashionMNIST(
        root="data",
        train=True,
        download=True,
        transform=ToTensor()
    )

    test_data = datasets.FashionMNIST(
        root="data",
        train=False,
        download=True,
        transform=ToTensor()
    )

    train_dataloader = DataLoader(training_data, batch_size=64)
    test_dataloader = DataLoader(test_data, batch_size=64)

    class NeuralNetwork(nn.Module):
        def __init__(self):
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear_relu_stack = nn.Sequential(
                nn.Linear(28*28, 512),
                nn.ReLU(),
                nn.Linear(512, 512),
                nn.ReLU(),
                nn.Linear(512, 10),
            )

        def forward(self, x):
            x = self.flatten(x)
            logits = self.linear_relu_stack(x)
            return logits

    model = NeuralNetwork()






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

 .. code-block:: none

    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz
    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz

      0%|          | 0/26421880 [00:00<?, ?it/s]
      0%|          | 65536/26421880 [00:00<01:12, 365712.47it/s]
      1%|          | 229376/26421880 [00:00<00:38, 686473.38it/s]
      4%|3         | 950272/26421880 [00:00<00:11, 2199132.79it/s]
     15%|#4        | 3833856/26421880 [00:00<00:02, 7640197.52it/s]
     33%|###2      | 8683520/26421880 [00:00<00:01, 17371306.45it/s]
     45%|####5     | 11894784/26421880 [00:00<00:00, 17867089.88it/s]
     68%|######7   | 17858560/26421880 [00:01<00:00, 23275184.21it/s]
     87%|########7 | 23003136/26421880 [00:01<00:00, 29472103.80it/s]
    100%|#########9| 26378240/26421880 [00:01<00:00, 26061594.96it/s]
    100%|##########| 26421880/26421880 [00:01<00:00, 18289954.84it/s]
    Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw

    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz
    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz

      0%|          | 0/29515 [00:00<?, ?it/s]
    100%|##########| 29515/29515 [00:00<00:00, 329254.84it/s]
    Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw

    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz
    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz

      0%|          | 0/4422102 [00:00<?, ?it/s]
      1%|1         | 65536/4422102 [00:00<00:11, 364558.11it/s]
      5%|5         | 229376/4422102 [00:00<00:06, 685149.97it/s]
     21%|##1       | 950272/4422102 [00:00<00:01, 2200190.39it/s]
     87%|########6 | 3833856/4422102 [00:00<00:00, 7648853.66it/s]
    100%|##########| 4422102/4422102 [00:00<00:00, 6119063.46it/s]
    Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw

    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz
    Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz

      0%|          | 0/5148 [00:00<?, ?it/s]
    100%|##########| 5148/5148 [00:00<00:00, 34603008.00it/s]
    Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw





.. GENERATED FROM PYTHON SOURCE LINES 71-83

Hyperparameters
-----------------

Hyperparameters are adjustable parameters that let you control the model optimization process.
Different hyperparameter values can impact model training and convergence rates
(`read more <https://pytorch.org/tutorials/beginner/hyperparameter_tuning_tutorial.html>`__ about hyperparameter tuning)

We define the following hyperparameters for training:
 - **Number of Epochs** - the number times to iterate over the dataset
 - **Batch Size** - the number of data samples propagated through the network before the parameters are updated
 - **Learning Rate** - how much to update models parameters at each batch/epoch. Smaller values yield slow learning speed, while large values may result in unpredictable behavior during training.


.. GENERATED FROM PYTHON SOURCE LINES 83-90

.. code-block:: default


    learning_rate = 1e-3
    batch_size = 64
    epochs = 5










.. GENERATED FROM PYTHON SOURCE LINES 91-117

Optimization Loop
-----------------

Once we set our hyperparameters, we can then train and optimize our model with an optimization loop. Each
iteration of the optimization loop is called an **epoch**.

Each epoch consists of two main parts:
 - **The Train Loop** - iterate over the training dataset and try to converge to optimal parameters.
 - **The Validation/Test Loop** - iterate over the test dataset to check if model performance is improving.

Let's briefly familiarize ourselves with some of the concepts used in the training loop. Jump ahead to
see the :ref:`full-impl-label` of the optimization loop.

Loss Function
~~~~~~~~~~~~~~~~~

When presented with some training data, our untrained network is likely not to give the correct
answer. **Loss function** measures the degree of dissimilarity of obtained result to the target value,
and it is the loss function that we want to minimize during training. To calculate the loss we make a
prediction using the inputs of our given data sample and compare it against the true data label value.

Common loss functions include `nn.MSELoss <https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html#torch.nn.MSELoss>`_ (Mean Square Error) for regression tasks, and
`nn.NLLLoss <https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html#torch.nn.NLLLoss>`_ (Negative Log Likelihood) for classification.
`nn.CrossEntropyLoss <https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss>`_ combines ``nn.LogSoftmax`` and ``nn.NLLLoss``.

We pass our model's output logits to ``nn.CrossEntropyLoss``, which will normalize the logits and compute the prediction error.

.. GENERATED FROM PYTHON SOURCE LINES 117-123

.. code-block:: default


    # Initialize the loss function
    loss_fn = nn.CrossEntropyLoss()










.. GENERATED FROM PYTHON SOURCE LINES 124-132

Optimizer
~~~~~~~~~~~~~~~~~

Optimization is the process of adjusting model parameters to reduce model error in each training step. **Optimization algorithms** define how this process is performed (in this example we use Stochastic Gradient Descent).
All optimization logic is encapsulated in  the ``optimizer`` object. Here, we use the SGD optimizer; additionally, there are many `different optimizers <https://pytorch.org/docs/stable/optim.html>`_
available in PyTorch such as ADAM and RMSProp, that work better for different kinds of models and data.

We initialize the optimizer by registering the model's parameters that need to be trained, and passing in the learning rate hyperparameter.

.. GENERATED FROM PYTHON SOURCE LINES 132-135

.. code-block:: default


    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)








.. GENERATED FROM PYTHON SOURCE LINES 136-140

Inside the training loop, optimization happens in three steps:
 * Call ``optimizer.zero_grad()`` to reset the gradients of model parameters. Gradients by default add up; to prevent double-counting, we explicitly zero them at each iteration.
 * Backpropagate the prediction loss with a call to ``loss.backward()``. PyTorch deposits the gradients of the loss w.r.t. each parameter.
 * Once we have our gradients, we call ``optimizer.step()`` to adjust the parameters by the gradients collected in the backward pass.

.. GENERATED FROM PYTHON SOURCE LINES 143-149

.. _full-impl-label:

Full Implementation
-----------------------
We define ``train_loop`` that loops over our optimization code, and ``test_loop`` that
evaluates the model's performance against our test data.

.. GENERATED FROM PYTHON SOURCE LINES 149-191

.. code-block:: default


    def train_loop(dataloader, model, loss_fn, optimizer):
        size = len(dataloader.dataset)
        # Set the model to training mode - important for batch normalization and dropout layers
        # Unnecessary in this situation but added for best practices
        model.train()
        for batch, (X, y) in enumerate(dataloader):
            # Compute prediction and loss
            pred = model(X)
            loss = loss_fn(pred, y)

            # Backpropagation
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            if batch % 100 == 0:
                loss, current = loss.item(), batch * batch_size + len(X)
                print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


    def test_loop(dataloader, model, loss_fn):
        # Set the model to evaluation mode - important for batch normalization and dropout layers
        # Unnecessary in this situation but added for best practices
        model.eval()
        size = len(dataloader.dataset)
        num_batches = len(dataloader)
        test_loss, correct = 0, 0

        # Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode
        # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True
        with torch.no_grad():
            for X, y in dataloader:
                pred = model(X)
                test_loss += loss_fn(pred, y).item()
                correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        test_loss /= num_batches
        correct /= size
        print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")









.. GENERATED FROM PYTHON SOURCE LINES 192-194

We initialize the loss function and optimizer, and pass it to ``train_loop`` and ``test_loop``.
Feel free to increase the number of epochs to track the model's improving performance.

.. GENERATED FROM PYTHON SOURCE LINES 194-207

.. code-block:: default


    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    epochs = 10
    for t in range(epochs):
        print(f"Epoch {t+1}\n-------------------------------")
        train_loop(train_dataloader, model, loss_fn, optimizer)
        test_loop(test_dataloader, model, loss_fn)
    print("Done!")







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

 .. code-block:: none

    Epoch 1
    -------------------------------
    loss: 2.298730  [   64/60000]
    loss: 2.289123  [ 6464/60000]
    loss: 2.273286  [12864/60000]
    loss: 2.269406  [19264/60000]
    loss: 2.249603  [25664/60000]
    loss: 2.229407  [32064/60000]
    loss: 2.227368  [38464/60000]
    loss: 2.204261  [44864/60000]
    loss: 2.206193  [51264/60000]
    loss: 2.166651  [57664/60000]
    Test Error: 
     Accuracy: 50.9%, Avg loss: 2.166725 

    Epoch 2
    -------------------------------
    loss: 2.176750  [   64/60000]
    loss: 2.169595  [ 6464/60000]
    loss: 2.117500  [12864/60000]
    loss: 2.129272  [19264/60000]
    loss: 2.079674  [25664/60000]
    loss: 2.032928  [32064/60000]
    loss: 2.050115  [38464/60000]
    loss: 1.985236  [44864/60000]
    loss: 1.987887  [51264/60000]
    loss: 1.907162  [57664/60000]
    Test Error: 
     Accuracy: 55.9%, Avg loss: 1.915486 

    Epoch 3
    -------------------------------
    loss: 1.951612  [   64/60000]
    loss: 1.928685  [ 6464/60000]
    loss: 1.815709  [12864/60000]
    loss: 1.841552  [19264/60000]
    loss: 1.732467  [25664/60000]
    loss: 1.692914  [32064/60000]
    loss: 1.701714  [38464/60000]
    loss: 1.610632  [44864/60000]
    loss: 1.632870  [51264/60000]
    loss: 1.514263  [57664/60000]
    Test Error: 
     Accuracy: 58.8%, Avg loss: 1.541525 

    Epoch 4
    -------------------------------
    loss: 1.616448  [   64/60000]
    loss: 1.582892  [ 6464/60000]
    loss: 1.427595  [12864/60000]
    loss: 1.487950  [19264/60000]
    loss: 1.359332  [25664/60000]
    loss: 1.364817  [32064/60000]
    loss: 1.371491  [38464/60000]
    loss: 1.298706  [44864/60000]
    loss: 1.336201  [51264/60000]
    loss: 1.232145  [57664/60000]
    Test Error: 
     Accuracy: 62.2%, Avg loss: 1.260237 

    Epoch 5
    -------------------------------
    loss: 1.345538  [   64/60000]
    loss: 1.327798  [ 6464/60000]
    loss: 1.153802  [12864/60000]
    loss: 1.254829  [19264/60000]
    loss: 1.117322  [25664/60000]
    loss: 1.153248  [32064/60000]
    loss: 1.171765  [38464/60000]
    loss: 1.110263  [44864/60000]
    loss: 1.154467  [51264/60000]
    loss: 1.070921  [57664/60000]
    Test Error: 
     Accuracy: 64.1%, Avg loss: 1.089831 

    Epoch 6
    -------------------------------
    loss: 1.166889  [   64/60000]
    loss: 1.170514  [ 6464/60000]
    loss: 0.979435  [12864/60000]
    loss: 1.113774  [19264/60000]
    loss: 0.973411  [25664/60000]
    loss: 1.015192  [32064/60000]
    loss: 1.051113  [38464/60000]
    loss: 0.993591  [44864/60000]
    loss: 1.039709  [51264/60000]
    loss: 0.971077  [57664/60000]
    Test Error: 
     Accuracy: 65.8%, Avg loss: 0.982440 

    Epoch 7
    -------------------------------
    loss: 1.045165  [   64/60000]
    loss: 1.070583  [ 6464/60000]
    loss: 0.862304  [12864/60000]
    loss: 1.022265  [19264/60000]
    loss: 0.885213  [25664/60000]
    loss: 0.919528  [32064/60000]
    loss: 0.972762  [38464/60000]
    loss: 0.918728  [44864/60000]
    loss: 0.961629  [51264/60000]
    loss: 0.904379  [57664/60000]
    Test Error: 
     Accuracy: 66.9%, Avg loss: 0.910167 

    Epoch 8
    -------------------------------
    loss: 0.956964  [   64/60000]
    loss: 1.002171  [ 6464/60000]
    loss: 0.779057  [12864/60000]
    loss: 0.958409  [19264/60000]
    loss: 0.827240  [25664/60000]
    loss: 0.850262  [32064/60000]
    loss: 0.917320  [38464/60000]
    loss: 0.868384  [44864/60000]
    loss: 0.905506  [51264/60000]
    loss: 0.856353  [57664/60000]
    Test Error: 
     Accuracy: 68.3%, Avg loss: 0.858248 

    Epoch 9
    -------------------------------
    loss: 0.889765  [   64/60000]
    loss: 0.951220  [ 6464/60000]
    loss: 0.717035  [12864/60000]
    loss: 0.911042  [19264/60000]
    loss: 0.786085  [25664/60000]
    loss: 0.798370  [32064/60000]
    loss: 0.874939  [38464/60000]
    loss: 0.832796  [44864/60000]
    loss: 0.863254  [51264/60000]
    loss: 0.819742  [57664/60000]
    Test Error: 
     Accuracy: 69.5%, Avg loss: 0.818780 

    Epoch 10
    -------------------------------
    loss: 0.836395  [   64/60000]
    loss: 0.910220  [ 6464/60000]
    loss: 0.668506  [12864/60000]
    loss: 0.874338  [19264/60000]
    loss: 0.754805  [25664/60000]
    loss: 0.758453  [32064/60000]
    loss: 0.840451  [38464/60000]
    loss: 0.806153  [44864/60000]
    loss: 0.830360  [51264/60000]
    loss: 0.790281  [57664/60000]
    Test Error: 
     Accuracy: 71.0%, Avg loss: 0.787271 

    Done!




.. GENERATED FROM PYTHON SOURCE LINES 208-214

Further Reading
-----------------------
- `Loss Functions <https://pytorch.org/docs/stable/nn.html#loss-functions>`_
- `torch.optim <https://pytorch.org/docs/stable/optim.html>`_
- `Warmstart Training a Model <https://pytorch.org/tutorials/recipes/recipes/warmstarting_model_using_parameters_from_a_different_model.html>`_



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

   **Total running time of the script:** ( 2 minutes  8.401 seconds)


.. _sphx_glr_download_beginner_basics_optimization_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: optimization_tutorial.py <optimization_tutorial.py>`

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

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


.. only:: html

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

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