{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 2\n",
    "\n",
    "In this tutorial, we'll look at gradient descent, its variants and optimization algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision import transforms\n",
    "\n",
    "print(\"Torch version:\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first load the MNIST dataset and split it into a train and test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train = datasets.MNIST(\n",
    "    root = \"data\",\n",
    "    train = True,\n",
    "    download = True,\n",
    "    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0, 1)])\n",
    ")\n",
    "data_test = datasets.MNIST(\n",
    "    root = \"data\",\n",
    "    train = False,\n",
    "    download = True,\n",
    "    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0, 1)])\n",
    ")\n",
    "\n",
    "dataloader_train = DataLoader(data_train, batch_size = batch_size, shuffle = True)\n",
    "dataloader_test = DataLoader(data_test, batch_size = batch_size, shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's defined helper functions totrain the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_accuracy(model, dataloader):\n",
    "    num_correct = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for (X, y) in dataloader:\n",
    "            pred = model(X)\n",
    "            num_correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    accuracy = num_correct / len(dataloader.dataset)\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_loss_accuracy(model, dataloader, loss_fn):\n",
    "    loss = 0\n",
    "    num_correct = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for (X, y) in dataloader:\n",
    "            pred = model(X)\n",
    "            loss += loss_fn(pred, y).item()\n",
    "            num_correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    loss /= len(dataloader)\n",
    "    accuracy = num_correct / len(dataloader.dataset)\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, loss_fn, optimizer, epochs, dataloader_train, dataloader_test, early_stopper = None, log_period = 10000):\n",
    "    for epoch in range(epochs):\n",
    "        processed_since_log = 0\n",
    "        for batch, (X, y) in enumerate(dataloader_train):\n",
    "            model.train()\n",
    "            pred = model(X)\n",
    "            loss = loss_fn(pred, y)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            processed_since_log += dataloader_train.batch_size\n",
    "\n",
    "            if processed_since_log >= log_period:\n",
    "                current = min((batch + 1) * dataloader_train.batch_size, len(data_train))\n",
    "                loss = loss.item()\n",
    "                model.eval()\n",
    "                train_acc = calculate_accuracy(model, dataloader_train)\n",
    "                test_loss, test_acc = calculate_loss_accuracy(model, dataloader_test, loss_fn)\n",
    "                print(f\"train loss: {loss:>7f}  test loss: {test_loss:>7f}  train accuracy: {train_acc:>3f}  test accuracy: {test_acc:>3f}  [sample {current:>5d}/{len(data_train):>5d}] [epoch {epoch+1:>2d}/{epochs:>2d}]\")\n",
    "                processed_since_log -= log_period"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we will be creating multiple identical models, let's also create a function for that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    model = nn.Sequential()\n",
    "    model.append(nn.Flatten())\n",
    "    model.append(nn.Linear(data_train.data.shape[1] * data_train.data.shape[2], 100))\n",
    "    model.append(nn.ReLU())\n",
    "    model.append(nn.Linear(100, 10))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Simple_SGD(torch.optim.Optimizer):\n",
    "    def __init__(self, params, lr = 1e-3):\n",
    "        defaults = dict(lr = lr)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad\n",
    "    def step(self):\n",
    "        for group in self.param_groups:\n",
    "            for p in group[\"params\"]:\n",
    "                p.add_(p.grad, alpha = -group[\"lr\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "optimizer = Simple_SGD(model.parameters())\n",
    "train_model(model, nn.CrossEntropyLoss(), optimizer, 10, dataloader_train, dataloader_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Simple_Momentum(torch.optim.Optimizer):\n",
    "    def __init__(self, params, lr = 1e-3, momentum_gamma = 0.9):\n",
    "        defaults = dict(lr = lr, momentum_gamma = momentum_gamma)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad\n",
    "    def step(self):\n",
    "        for group in self.param_groups:\n",
    "            for p in group[\"params\"]:\n",
    "                momentum_v = self.state[p].get(\"momentum_v\")\n",
    "                if momentum_v is None:\n",
    "                    momentum_v = torch.clone(p.grad).detach()\n",
    "                else:\n",
    "                    momentum_v.mul_(group[\"momentum_gamma\"]).add_(p.grad)\n",
    "                \n",
    "                self.state[p][\"momentum_v\"] = momentum_v\n",
    "                p.add_(momentum_v, alpha = -group[\"lr\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "optimizer = Simple_Momentum(model.parameters())\n",
    "train_model(model, nn.CrossEntropyLoss(), optimizer, 10, dataloader_train, dataloader_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1\n",
    "\n",
    "Create an optimizer for the Nesterov accelerated gradients (NAG) method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Simple_Nesterov(torch.optim.Optimizer):\n",
    "    def __init__(self, params, lr = 1e-3, momentum_gamma = 0.9):\n",
    "        defaults = dict(lr = lr, momentum_gamma = momentum_gamma)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad\n",
    "    def step(self):\n",
    "        for group in self.param_groups:\n",
    "            for p in group[\"params\"]:\n",
    "                # TODO\n",
    "                pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "optimizer = Simple_Nesterov(model.parameters())\n",
    "train_model(model, nn.CrossEntropyLoss(), optimizer, 10, dataloader_train, dataloader_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2\n",
    "\n",
    "Create an optimizer for the ADAM method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Simple_ADAM(torch.optim.Optimizer):\n",
    "    def __init__(self, params, lr = 1e-3, beta_1 = 0.9, beta_2 = 0.999, epsilon = 1e-8):\n",
    "        defaults = dict(lr = lr, beta_1 = beta_1, beta_2 = beta_2, epsilon = epsilon)\n",
    "        super().__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad\n",
    "    def step(self):\n",
    "        for group in self.param_groups:\n",
    "            for p in group[\"params\"]:\n",
    "                # TODO\n",
    "                pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "optimizer = Simple_ADAM(model.parameters())\n",
    "train_model(model, nn.CrossEntropyLoss(), optimizer, 10, dataloader_train, dataloader_test)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tzn",
   "language": "python",
   "name": "tzn"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}