pytorch/about_the_optimization_loop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing the model parameters\n",
    "\n",
    "Now that we have a model and it's time to train, validate and test our model by optimizing its parameters on your data. Training a model is an iterative process; in each iteration (called an *epoch*). 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 module), and **optimizes** these parameters using gradient descent.\n",
    "\n",
    "### Prerequisite code\n",
    "\n",
    "We will load the code from the previous articles [**Datasets & DataLoaders**](https://www.kaggle.com/code/aisuko/loading-and-normalizing-datasets) and [**Build Model**](https://www.kaggle.com/code/aisuko/building-the-model-layers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the paltform, Apple Silicon or Linux\n",
    "import os, platform\n",
    "\n",
    "torch_device=\"cpu\"\n",
    "\n",
    "if 'kaggle' in os.environ.get('KAGGLE_URL_BASE','localhost'):\n",
    "    torch_device = 'cuda'\n",
    "else:\n",
    "    torch_device = 'mps' if platform.system() == 'Darwin' else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor, Lambda\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64)\n",
    "test_dataloader = DataLoader(test_data, batch_size=64)\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(torch_device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting hyperparameters\n",
    "\n",
    "Hyperparameters are adjustable parameters that let you control the model optimization peocess. Different hyperparameter values can impact model training and the level of accuracy.\n",
    "\n",
    "We define the following hyperparameters for training:\n",
    "* **Number of Epochs** - the number times the entire training dataset is pass through the network.\n",
    "* **Batch Size** - the number of data samples seen by the model in each epoch. Iterates are the number of batches needs to complete an epoch.\n",
    "* **Learning Rate** - the size of steps the model match as it searchs for best weights that will produce a higher model accuracy. Smaller values mean the model will take a longer time to find the best weights, while larger values may result in the model step over and misses the best weights which yields unpredictable behavior during training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 64\n",
    "epochs = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add an optimization loop\n",
    "\n",
    "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**.\n",
    "\n",
    "Each epoch consists of two main parts:\n",
    "* **The Train Loop** - iterate over the training dataset and try to converge to optimal parameters.\n",
    "* **The Validation/Test Loop** - iterate over the best dataset to check if model performance is improving.\n",
    "\n",
    "Let's briefly familiarize ourselves with some of the concepts used in the training loop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add a loss function\n",
    "\n",
    "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.\n",
    "\n",
    "Common loss functions include:\n",
    "* `nn.MSELoss` (Mean Square Error) used for regression tasks\n",
    "* `nn.NLLLoss` (Negative Log Likelihood) used for classification\n",
    "* `nn.CrossEntropyLoss` combines `nn.LogSoftmax` and `nn.NLLLoss`\n",
    "\n",
    "We pass our model's output logits to `nn.CrossEntropyLoss`, which will normalize the logits and compute the prediction error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the loss function\n",
    "loss_fn = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimization pass\n",
    "\n",
    "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 available in PyTorch such as `ADAM and RMSProp`, that work better for different kinds of models and data.\n",
    "\n",
    "We initialize the optimizer by registering the model's parameters that need to be trained, and passing in the learning rate hyperparameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inside the training loop, optimization happens in three steps:\n",
    "* 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.\n",
    "* Back-propagate the prediction loss with a call to `loss.backwards()`. PyTorch deposits the gradients of the loss w.r.t each parameter.\n",
    "* Once we have our gradients, we call `optimizer.step()` to adjust the parameters by the gradients collected in the backward pass."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full implementation\n",
    "\n",
    "We define `train_loop` that loops over our optimization code, and `test_loop` that evaluates the model's performance against out test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Compute prediction and loss\n",
    "        # X input tensor\n",
    "        # y is a tensor of target labels\n",
    "        X, y = X.to(torch_device), y.to(torch_device)\n",
    "        pred = model (X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if batch % 100 ==0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f} [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    test_loss, correct = 0,0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for X,y in dataloader:\n",
    "            X, y = X.to(torch_device), y.to(torch_device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1)==y).type(torch.float).sum().item()\n",
    "    test_loss /= size\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "epochs = 10\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
    "    test_loop(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try running the loop for more `epochs` or adjusting the `learning_rate` to a bigger number. It might also be the case that the model configuration we choise might not be the optimal one for this kind of problem (it isn't)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving models\n",
    "\n",
    "When you are satisfied with the model's performance, you can use `torch.save` to save it. PyTorch models store the learned parameters in an internal state dictionary, called `state_dict`. These can be persisted with the `torch.save` method:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"model.pth\")\n",
    "print(\"Saved PyTorch Model State to model.pth\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "imp",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}