conv2d.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import warnings\n",
    "import dataclasses as dc\n",
    "import typing as tp\n",
    "warnings.filterwarnings(\"ignore\") # W0901 12:56:55.922000 133054240231424 torch/fx/experimental/symbolic_shapes.py:4449] [0/1] xindex is not in var_ranges, defaulting to unknown range.\n",
    "\n",
    "import torch\n",
    "from torch import Tensor, nn, _assert\n",
    "from torch.nn import functional as F\n",
    "import torch.utils.data\n",
    "from torchvision import datasets, transforms\n",
    "DEVICE = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DataLoading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader:\n",
    "    def __init__(self, train:bool):\n",
    "        transform = transforms.Compose([\n",
    "            transforms.ToTensor(), # (H, W, C)/(H, W) -> (C, H, W) AND [0, 255] -> [0.0, 1.0]\n",
    "        ])\n",
    "        if train:\n",
    "            self.ds = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "        else:\n",
    "            self.ds = datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "        \n",
    "    def iter_batches(self, batch_size):\n",
    "        while True:\n",
    "            self.dataset = torch.utils.data.DataLoader(\n",
    "                dataset=self.ds,\n",
    "                batch_size=batch_size,\n",
    "                shuffle=True,\n",
    "                pin_memory=True,\n",
    "                drop_last=True\n",
    "            )\n",
    "            for X_batch, y_batch in self.dataset:\n",
    "                yield X_batch.to(DEVICE), y_batch.to(DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LeNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "FAN_IN:tp.TypeAlias = int\n",
    "FAN_OUT:tp.TypeAlias = int\n",
    "KERNEL_SIZE:tp.TypeAlias = tuple[int, int]\n",
    "STRIDES:tp.TypeAlias = tuple[int, int]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dc.dataclass\n",
    "class Wei:\n",
    "    wei:Tensor\n",
    "    grad:tp.Optional[Tensor] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(\n",
    "    kernel1:tuple[FAN_OUT, FAN_IN, KERNEL_SIZE] = (6, 1, (5, 5)),\n",
    "    kernel2:tuple[FAN_OUT, FAN_IN, KERNEL_SIZE] = (16, 6, (5, 5)),\n",
    "    kernel3:tuple[FAN_OUT, FAN_IN, KERNEL_SIZE] = (120, 16, (5, 5)),\n",
    "    weight4:tuple[FAN_OUT, FAN_IN] = (84, 120),\n",
    "    weight5:tuple[FAN_OUT, FAN_IN] = (10, 84)\n",
    "):\n",
    "    sqrt2 = 2**0.5\n",
    "    prod = lambda x: x[0]*x[1]\n",
    "\n",
    "    # First Convolutional Layer: Stride 1\n",
    "    bound = sqrt2/prod(kernel1[-1])\n",
    "    W1 = Wei(torch.empty(size=(kernel1[0], kernel1[1], *kernel1[-1]), requires_grad=True).uniform_(-bound, bound))\n",
    "    B1 = Wei(torch.zeros(size=kernel1[0], requires_grad=True))\n",
    "\n",
    "    # Max Pooling Layer: Stride 2, Kernel 2x2, No Weights\n",
    "\n",
    "    # Second Convolutional Layer: Stride 1\n",
    "    bound = sqrt2/prod(kernel2[-1])\n",
    "    W2 = Wei(torch.empty(size=(kernel2[0], kernel2[1], *kernel2[-1]), requires_grad=True).uniform_(-bound, bound))\n",
    "    B2 = Wei(torch.zeros(size=kernel2[0], requires_grad=True))\n",
    "\n",
    "    # Max Pooling Layer: Stride 2, Kernel 2x2, No Weights\n",
    "\n",
    "    # Third Convolutional Layer\n",
    "    bound = sqrt2/prod(kernel3[-1])\n",
    "    W3 = Wei(torch.empty(size=(kernel3[0], kernel3[1], *kernel3[-1]), requires_grad=True).uniform_(-bound, bound))\n",
    "    B3 = Wei(torch.zeros(size=kernel3[0], requires_grad=True))\n",
    "\n",
    "    # First Linear Layer\n",
    "    bound = sqrt2/prod(weight4[-1])\n",
    "    W4 = Wei(torch.empty(size=weight4, requires_grad=True).uniform_(-bound, bound))\n",
    "    B4 = Wei(torch.zeros(size=weight4[0], requires_grad=True))\n",
    "\n",
    "    # Second Linear Layer\n",
    "    bound = sqrt2/prod(weight5[-1])\n",
    "    W5 = Wei(torch.empty(size=weight5, requires_grad=True).uniform_(-bound, bound))\n",
    "    B5 = Wei(torch.zeros(size=weight5[0], requires_grad=True))\n",
    "\n",
    "    return W1, B1, W2, B2, W3, B3, W4, B4, W5, B5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conv2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _conv2d(\n",
    "    x:Tensor, # (H, W)\n",
    "    w:Tensor, # (h, w)\n",
    "    full:bool=False,\n",
    "    convolve:bool=False\n",
    "):\n",
    "    x = x[None, None, ...] # (1, 1, H, W)\n",
    "    w = w[None, None, ...] # (1, 1, h, w)\n",
    "    if full:\n",
    "        pad_h = w.size(-2) - 1\n",
    "        pad_w = w.size(-1) - 1\n",
    "        x = F.pad(x, (pad_w, pad_w, pad_h, pad_h), mode='constant')\n",
    "    if convolve:\n",
    "        w = w.flip((2, 3))\n",
    "    return F.conv2d(\n",
    "        input=x, # (1, 1, H, W)\n",
    "        weight=w, # (1, 1, h, w)\n",
    "        stride=1\n",
    "    )[0, 0, :, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from scipy.signal import correlate2d, convolve2d\n",
    "x = torch.randn(size=(2, 1, 28, 28), requires_grad=True)\n",
    "w = torch.randn(size=(6, 1, 5, 5), requires_grad=True)\n",
    "b = torch.randn(size=(6,), requires_grad=True)\n",
    "\n",
    "# \"full\" testing\n",
    "ful = correlate2d(x[0, 0].detach().numpy(), w[0, 0].detach().numpy(), mode='full')\n",
    "myful = _conv2d(x[0, 0].detach(), w[0, 0].detach(), full=True)\n",
    "diff_full = (ful-myful.numpy())\n",
    "diff_full.mean(), diff_full.std(), abs(diff_full).max(), abs(diff_full).min()\n",
    "\n",
    "# \"convolve\" testing\n",
    "my_convolve = _conv2d(x[0, 0].detach(), w[0, 0].detach(), convolve=True, full=True)\n",
    "convolve = convolve2d(x[0, 0].detach().numpy(), w[0, 0].detach().numpy(), mode='full')\n",
    "diff_convolve = (convolve-my_convolve.numpy())\n",
    "diff_convolve.mean(), diff_convolve.std(), abs(diff_convolve).max(), abs(diff_convolve).min()\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv2d_forward(\n",
    "    x:Tensor,               # (B, fi, H, W) \n",
    "    wie:Tensor,               # (fo, fi, h, w)\n",
    "    bias:tp.Optional[Tensor],  # (fo,)\n",
    "    stride:STRIDES = (1, 1)\n",
    "):\n",
    "    B, C, H, W = x.shape\n",
    "    fo, fi, h, w = wie.shape\n",
    "    sh, sw = stride\n",
    "    assert C == fi, f\"Expected {C} == {fi}\"\n",
    "    assert H >= h, f\"Expected {H} >= {h}\"\n",
    "    assert W >= w, f\"Expected {W} >= {w}\"\n",
    "    assert bias.shape[0] == fo, f\"Expected {bias.shape[0]} == {fo}\"\n",
    "\n",
    "    output_shape = (B, fo, int((H-h)//sh + 1), int((W-w)//sw + 1))\n",
    "    O = torch.zeros(output_shape) # (B, C_out, H1, W1)\n",
    "    for fan_out in range(fo):\n",
    "        for fan_in in range(fi):\n",
    "            for bdim in range(B):\n",
    "                O[bdim, fan_out] += _conv2d(x[bdim, fan_in], wie[fan_out, fan_in])[::sh, ::sw]\n",
    "\n",
    "    if bias is not None:\n",
    "        O += bias.view(1, -1, 1, 1)\n",
    "    return O"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _dilate_matrix(x:Tensor, dilation:tuple[int, int]):\n",
    "    \"\"\"`x: shape(B, C, H, W)`\\n `dilation:tuple`\"\"\"\n",
    "    (B, C, H, W), (Hd, Wd)  = x.shape, dilation\n",
    "    dilated = torch.zeros((B, C, Hd*(H-1)+1, Wd*(W-1)+1 ))\n",
    "    dilated[:, :, ::Hd, ::Wd] = x\n",
    "    return dilated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv2d_backward(\n",
    "    x:Tensor,                   # (B, fi, H, W)\n",
    "    wei:Tensor,                 # (fo, fi, h, w)\n",
    "    bias:tp.Optional[Tensor],   # (fo,)\n",
    "    dL_dO:Tensor,               # (B, fo, H1, W1)\n",
    "    stride:STRIDES = (1, 1)\n",
    "):  \n",
    "    fo, fi, h, w = wei.shape\n",
    "    B, C, H, W = x.shape\n",
    "\n",
    "    dL_dx, dL_dwei = torch.zeros_like(x), torch.zeros_like(wei)\n",
    "    dL_dO = _dilate_matrix(dL_dO, dilation=stride)\n",
    "    for fan_out in range(fo): # C_out\n",
    "        for bdim in range(B): # B\n",
    "            for fan_in in range(fi): # C_in\n",
    "                dL_dwei[fan_out, fan_in] += _conv2d(x[bdim, fan_in], dL_dO[bdim, fan_out]) # (H, W)*(H1, W1) => (Hk, Wk)\n",
    "                dL_dx[bdim, fan_in] += _conv2d(dL_dO[bdim, fan_out], wei[fan_out, fan_in], full=True, convolve=True) # (H1, W1)*(Hk, Wk) => (H, W)\n",
    "    \n",
    "    dL_db = dL_dO.sum(dim=(0, 2, 3)) if bias is not None else None\n",
    "    return dL_dx, dL_dwei, dL_db"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test conv2d forward and backward methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x = torch.randn(size=(2, 1, 28, 28), requires_grad=True)\n",
    "# w = torch.randn(size=(6, 1, 5, 5), requires_grad=True)\n",
    "# b = torch.randn(size=(6,), requires_grad=True)\n",
    "\n",
    "# my_conv2d = conv2d_forward(x, w, b)\n",
    "# my_conv2d.retain_grad()\n",
    "\n",
    "# torch_conv2d = F.conv2d(x, w, b, stride=1)\n",
    "# torch.testing.assert_close(my_conv2d.detach(), torch_conv2d)\n",
    "\n",
    "# loss = my_conv2d.mean()\n",
    "# loss.backward()\n",
    "\n",
    "# torch_dL_dx, torch_dL_dw, torch_dL_db = x.grad, w.grad, b.grad\n",
    "# dL_dO = my_conv2d.grad.clone()\n",
    "\n",
    "# my_dL_dx, my_dL_dw, my_dL_db = conv2d_backward(x, w, b, dL_dO)\n",
    "# torch.testing.assert_close(my_dL_dx, torch_dL_dx)\n",
    "# torch.testing.assert_close(my_dL_dw, torch_dL_dw)\n",
    "# torch.testing.assert_close(my_dL_db, torch_dL_db)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MaxPool2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _maxpool(matrix:Tensor, kernel_size:KERNEL_SIZE, strides:STRIDES): # (H, W)\n",
    "        (H, W), (Hk, Wk), (Hs, Ws) = matrix.shape, kernel_size, strides\n",
    "        output_shape = ((H-Hk+1)//Hs + 1, (W-Wk+1)//Ws + 1)\n",
    "        indices, maxpooled = [], []\n",
    "        for i in range(0, H - Hk + 1, Hs):\n",
    "            for j in range(0, W - Wk + 1, Ws):\n",
    "                window = matrix[i:i+Hk, j:j+Wk]\n",
    "                max_index = torch.unravel_index(torch.argmax(window), window.shape)\n",
    "                max_index_global = (max_index[0] + i, max_index[1] + j)\n",
    "                indices.append(max_index_global)\n",
    "                maxpooled.append(window[max_index])\n",
    "        maxpooled = torch.tensor(maxpooled).reshape(output_shape)\n",
    "        indices = torch.tensor(indices) # (H1*W1, 2)\n",
    "        # (H1, W1), ((H1, W1), (H1, W1))\n",
    "        return maxpooled, (indices[:, 0].reshape(output_shape), indices[:, 1].reshape(output_shape))\n",
    "\n",
    "channeled_maxpool = torch.vmap(_maxpool, in_dims=(0, None, None), out_dims=0) # (C, H, W)\n",
    "vmaxpool = torch.vmap(channeled_maxpool, in_dims=(0, None, None), out_dims=0) # (B, C, H, W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def channeled_maxpool(matrix:Tensor, kernel_size:KERNEL_SIZE, strides:STRIDES): # (C, H, W)\n",
    "    (C, H, W) = matrix.shape\n",
    "    maxpooled, Rindices, Cindices = [], [], []\n",
    "    for c in range(C):\n",
    "        maxpooled_, (indices_r, indices_c) = _maxpool(matrix[c], kernel_size, strides)\n",
    "        maxpooled.append(maxpooled_)\n",
    "        Rindices.append(indices_r)\n",
    "        Cindices.append(indices_c)\n",
    "    return torch.stack(maxpooled), (torch.stack(Rindices), torch.stack(Cindices))\n",
    "\n",
    "def vmaxpool(matrix:Tensor, kernel_size:KERNEL_SIZE, strides:STRIDES): # (B, C, H, W)\n",
    "    (B, C, H, W) = matrix.shape\n",
    "    maxpooled, Rindices, Cindices = [], [], []\n",
    "    for b in range(B):\n",
    "        maxpooled_b, (indices_r, indices_c) = channeled_maxpool(matrix[b], kernel_size, strides)\n",
    "        maxpooled.append(maxpooled_b)\n",
    "        Rindices.append(indices_r)\n",
    "        Cindices.append(indices_c)\n",
    "    return torch.stack(maxpooled), (torch.stack(Rindices), torch.stack(Cindices))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def maxpool2d_forward(\n",
    "    x:Tensor,\n",
    "    kernel_size:KERNEL_SIZE,\n",
    "    strides:STRIDES,\n",
    ") -> tuple[Tensor, tuple[Tensor, Tensor], torch.Size]:\n",
    "    O, (ridx, cidx) = vmaxpool(x, kernel_size, strides)\n",
    "    return O, (ridx, cidx), x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def maxpool_backward(\n",
    "    dL_dO:Tensor,\n",
    "    x_shape:torch.Size,\n",
    "    indices:tuple[Tensor, Tensor],\n",
    "):\n",
    "    \"\"\"SOMETHING IS WRONG\"\"\"\n",
    "    (ridx, cidx) = indices\n",
    "\n",
    "    dL_dY = torch.zeros(x_shape) # (B, C, H, W)\n",
    "    dL_dY[:, :, ridx, cidx] += dL_dO\n",
    "\n",
    "    dY_dX = torch.zeros(x_shape) # (B, C, H, W)\n",
    "    dY_dX[:, :, ridx, cidx] += 1.0\n",
    "\n",
    "    dL_dX = dL_dY*dY_dX # (B, C, H, W)\n",
    "    return dL_dX # (B, C, H, W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test maxpool2d forward and backward methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(size=(2, 1, 28, 28), requires_grad=True)\n",
    "x.requires_grad = True\n",
    "\n",
    "torch_maxpool, torch_idx = F.max_pool2d_with_indices(x, kernel_size=(2, 2), stride=(2, 2))\n",
    "torch_maxpool.retain_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_loss = torch_maxpool.mean()\n",
    "torch_loss.backward()\n",
    "\n",
    "torch_dL_dX, torch_dL_dO = x.grad, torch_maxpool.grad.clone()\n",
    "assert not any([torch_dL_dX is None, torch_dL_dO is None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-0.0005), tensor(0.0010), tensor(0.0026), tensor(0.))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_maxpool, indices, x_shape = maxpool2d_forward(x.clone(), kernel_size=(2, 2), strides=(2, 2))\n",
    "torch.testing.assert_close(my_maxpool, torch_maxpool)\n",
    "my_dL_dx = maxpool_backward(torch_dL_dO, x_shape, indices)\n",
    "\n",
    "diff = torch_dL_dX - my_dL_dx\n",
    "diff.mean(), diff.std(), diff.abs().max(), diff.abs().min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 29,   3,   5,  34,   9,  38,  41,  43,  45,  18,  21,  22,  52,  54],\n",
       "          [ 85,  58,  89,  91,  64,  67,  97,  98,  72, 102,  76,  79,  81,  83],\n",
       "          [141, 143, 145, 119, 120, 151, 125, 126, 129, 159, 161, 162, 165, 138],\n",
       "          [169, 171, 200, 175, 205, 206, 209, 211, 184, 215, 216, 219, 220, 195],\n",
       "          [252, 254, 228, 230, 232, 262, 265, 239, 241, 270, 273, 274, 277, 251],\n",
       "          [280, 283, 284, 315, 316, 319, 292, 322, 296, 326, 300, 330, 333, 306],\n",
       "          [364, 339, 369, 371, 373, 347, 376, 378, 380, 383, 385, 359, 388, 391],\n",
       "          [392, 422, 397, 426, 429, 403, 405, 407, 436, 411, 412, 443, 444, 418],\n",
       "          [449, 479, 480, 454, 457, 486, 488, 491, 493, 495, 468, 470, 500, 502],\n",
       "          [532, 535, 508, 539, 513, 515, 544, 546, 549, 550, 553, 555, 528, 530],\n",
       "          [561, 591, 564, 594, 596, 570, 601, 602, 604, 606, 609, 611, 585, 615],\n",
       "          [616, 618, 648, 651, 652, 655, 657, 659, 661, 662, 636, 638, 640, 670],\n",
       "          [673, 703, 705, 707, 681, 683, 712, 686, 717, 690, 693, 695, 725, 727],\n",
       "          [728, 759, 761, 763, 737, 739, 740, 743, 745, 746, 749, 751, 780, 783]]],\n",
       "\n",
       "\n",
       "        [[[  1,  30,   4,   6,   8,  10,  41,  43,  45,  46,  20,  50,  52,  27],\n",
       "          [ 85,  87,  61,  90,  93,  67,  97,  71, 100, 103, 104,  78, 108, 111],\n",
       "          [141, 114, 145, 147, 149, 122, 153, 154, 157, 159, 160, 135, 137, 139],\n",
       "          [169, 198, 173, 174, 205, 179, 181, 182, 213, 187, 188, 218, 220, 222],\n",
       "          [252, 255, 228, 258, 261, 263, 265, 238, 269, 243, 272, 274, 276, 279],\n",
       "          [280, 310, 313, 315, 288, 291, 321, 294, 325, 299, 300, 330, 304, 334],\n",
       "          [337, 366, 341, 371, 372, 347, 349, 378, 352, 354, 384, 358, 361, 390],\n",
       "          [421, 395, 424, 426, 428, 403, 432, 434, 436, 439, 441, 443, 416, 446],\n",
       "          [476, 478, 480, 483, 456, 486, 488, 463, 492, 467, 496, 470, 501, 502],\n",
       "          [504, 507, 508, 538, 540, 543, 517, 546, 549, 550, 553, 554, 557, 531],\n",
       "          [588, 591, 564, 594, 569, 598, 572, 574, 605, 579, 580, 611, 585, 615],\n",
       "          [617, 619, 620, 622, 625, 627, 629, 630, 633, 662, 664, 666, 668, 670],\n",
       "          [701, 674, 704, 706, 709, 710, 713, 715, 717, 718, 693, 694, 724, 726],\n",
       "          [757, 730, 732, 762, 737, 738, 741, 770, 745, 775, 749, 751, 752, 782]]]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[[ 1,  0,  0,  1,  0,  1,  1,  1,  1,  0,  0,  0,  1,  1],\n",
       "           [ 3,  2,  3,  3,  2,  2,  3,  3,  2,  3,  2,  2,  2,  2],\n",
       "           [ 5,  5,  5,  4,  4,  5,  4,  4,  4,  5,  5,  5,  5,  4],\n",
       "           [ 6,  6,  7,  6,  7,  7,  7,  7,  6,  7,  7,  7,  7,  6],\n",
       "           [ 9,  9,  8,  8,  8,  9,  9,  8,  8,  9,  9,  9,  9,  8],\n",
       "           [10, 10, 10, 11, 11, 11, 10, 11, 10, 11, 10, 11, 11, 10],\n",
       "           [13, 12, 13, 13, 13, 12, 13, 13, 13, 13, 13, 12, 13, 13],\n",
       "           [14, 15, 14, 15, 15, 14, 14, 14, 15, 14, 14, 15, 15, 14],\n",
       "           [16, 17, 17, 16, 16, 17, 17, 17, 17, 17, 16, 16, 17, 17],\n",
       "           [19, 19, 18, 19, 18, 18, 19, 19, 19, 19, 19, 19, 18, 18],\n",
       "           [20, 21, 20, 21, 21, 20, 21, 21, 21, 21, 21, 21, 20, 21],\n",
       "           [22, 22, 23, 23, 23, 23, 23, 23, 23, 23, 22, 22, 22, 23],\n",
       "           [24, 25, 25, 25, 24, 24, 25, 24, 25, 24, 24, 24, 25, 25],\n",
       "           [26, 27, 27, 27, 26, 26, 26, 26, 26, 26, 26, 26, 27, 27]]],\n",
       " \n",
       " \n",
       "         [[[ 0,  1,  0,  0,  0,  0,  1,  1,  1,  1,  0,  1,  1,  0],\n",
       "           [ 3,  3,  2,  3,  3,  2,  3,  2,  3,  3,  3,  2,  3,  3],\n",
       "           [ 5,  4,  5,  5,  5,  4,  5,  5,  5,  5,  5,  4,  4,  4],\n",
       "           [ 6,  7,  6,  6,  7,  6,  6,  6,  7,  6,  6,  7,  7,  7],\n",
       "           [ 9,  9,  8,  9,  9,  9,  9,  8,  9,  8,  9,  9,  9,  9],\n",
       "           [10, 11, 11, 11, 10, 10, 11, 10, 11, 10, 10, 11, 10, 11],\n",
       "           [12, 13, 12, 13, 13, 12, 12, 13, 12, 12, 13, 12, 12, 13],\n",
       "           [15, 14, 15, 15, 15, 14, 15, 15, 15, 15, 15, 15, 14, 15],\n",
       "           [17, 17, 17, 17, 16, 17, 17, 16, 17, 16, 17, 16, 17, 17],\n",
       "           [18, 18, 18, 19, 19, 19, 18, 19, 19, 19, 19, 19, 19, 18],\n",
       "           [21, 21, 20, 21, 20, 21, 20, 20, 21, 20, 20, 21, 20, 21],\n",
       "           [22, 22, 22, 22, 22, 22, 22, 22, 22, 23, 23, 23, 23, 23],\n",
       "           [25, 24, 25, 25, 25, 25, 25, 25, 25, 25, 24, 24, 25, 25],\n",
       "           [27, 26, 26, 27, 26, 26, 26, 27, 26, 27, 26, 26, 26, 27]]]]),\n",
       " tensor([[[[ 1,  3,  5,  6,  9, 10, 13, 15, 17, 18, 21, 22, 24, 26],\n",
       "           [ 1,  2,  5,  7,  8, 11, 13, 14, 16, 18, 20, 23, 25, 27],\n",
       "           [ 1,  3,  5,  7,  8, 11, 13, 14, 17, 19, 21, 22, 25, 26],\n",
       "           [ 1,  3,  4,  7,  9, 10, 13, 15, 16, 19, 20, 23, 24, 27],\n",
       "           [ 0,  2,  4,  6,  8, 10, 13, 15, 17, 18, 21, 22, 25, 27],\n",
       "           [ 0,  3,  4,  7,  8, 11, 12, 14, 16, 18, 20, 22, 25, 26],\n",
       "           [ 0,  3,  5,  7,  9, 11, 12, 14, 16, 19, 21, 23, 24, 27],\n",
       "           [ 0,  2,  5,  6,  9, 11, 13, 15, 16, 19, 20, 23, 24, 26],\n",
       "           [ 1,  3,  4,  6,  9, 10, 12, 15, 17, 19, 20, 22, 24, 26],\n",
       "           [ 0,  3,  4,  7,  9, 11, 12, 14, 17, 18, 21, 23, 24, 26],\n",
       "           [ 1,  3,  4,  6,  8, 10, 13, 14, 16, 18, 21, 23, 25, 27],\n",
       "           [ 0,  2,  4,  7,  8, 11, 13, 15, 17, 18, 20, 22, 24, 26],\n",
       "           [ 1,  3,  5,  7,  9, 11, 12, 14, 17, 18, 21, 23, 25, 27],\n",
       "           [ 0,  3,  5,  7,  9, 11, 12, 15, 17, 18, 21, 23, 24, 27]]],\n",
       " \n",
       " \n",
       "         [[[ 1,  2,  4,  6,  8, 10, 13, 15, 17, 18, 20, 22, 24, 27],\n",
       "           [ 1,  3,  5,  6,  9, 11, 13, 15, 16, 19, 20, 22, 24, 27],\n",
       "           [ 1,  2,  5,  7,  9, 10, 13, 14, 17, 19, 20, 23, 25, 27],\n",
       "           [ 1,  2,  5,  6,  9, 11, 13, 14, 17, 19, 20, 22, 24, 26],\n",
       "           [ 0,  3,  4,  6,  9, 11, 13, 14, 17, 19, 20, 22, 24, 27],\n",
       "           [ 0,  2,  5,  7,  8, 11, 13, 14, 17, 19, 20, 22, 24, 26],\n",
       "           [ 1,  2,  5,  7,  8, 11, 13, 14, 16, 18, 20, 22, 25, 26],\n",
       "           [ 1,  3,  4,  6,  8, 11, 12, 14, 16, 19, 21, 23, 24, 26],\n",
       "           [ 0,  2,  4,  7,  8, 10, 12, 15, 16, 19, 20, 22, 25, 26],\n",
       "           [ 0,  3,  4,  6,  8, 11, 13, 14, 17, 18, 21, 22, 25, 27],\n",
       "           [ 0,  3,  4,  6,  9, 10, 12, 14, 17, 19, 20, 23, 25, 27],\n",
       "           [ 1,  3,  4,  6,  9, 11, 13, 14, 17, 18, 20, 22, 24, 26],\n",
       "           [ 1,  2,  4,  6,  9, 10, 13, 15, 17, 18, 21, 22, 24, 26],\n",
       "           [ 1,  2,  4,  6,  9, 10, 13, 14, 17, 19, 21, 23, 24, 26]]]]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[[0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0026, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0026,  ..., 0.0026, 0.0000, 0.0026],\n",
       "           ...,\n",
       "           [0.0000, 0.0000, 0.0000,  ..., 0.0026, 0.0000, 0.0026],\n",
       "           [0.0026, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0026]]],\n",
       " \n",
       " \n",
       "         [[[0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0000, 0.0026],\n",
       "           [0.0000, 0.0000, 0.0026,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           ...,\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0026, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0026,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0026, 0.0000]]]]),\n",
       " torch.Size([2, 1, 28, 28]))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_dL_dX, torch_dL_dX.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[[0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0000, 0.0026],\n",
       "           [0.0000, 0.0026, 0.0026,  ..., 0.0000, 0.0026, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0026,  ..., 0.0026, 0.0000, 0.0026],\n",
       "           ...,\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0026, 0.0026, 0.0026],\n",
       "           [0.0026, 0.0000, 0.0026,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0026, 0.0026]]],\n",
       " \n",
       " \n",
       "         [[[0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0000, 0.0026],\n",
       "           [0.0000, 0.0026, 0.0026,  ..., 0.0000, 0.0026, 0.0000],\n",
       "           [0.0000, 0.0000, 0.0026,  ..., 0.0026, 0.0000, 0.0026],\n",
       "           ...,\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0026, 0.0026, 0.0026],\n",
       "           [0.0026, 0.0000, 0.0026,  ..., 0.0000, 0.0000, 0.0000],\n",
       "           [0.0000, 0.0026, 0.0000,  ..., 0.0000, 0.0026, 0.0026]]]]),\n",
       " torch.Size([2, 1, 28, 28]))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_dL_dx, my_dL_dx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(1280), 1568)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(my_dL_dx == torch_dL_dX).sum(), torch_dL_dX.numel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.7207, -0.7193, -0.4922,  ...,  0.4852,  0.5819, -0.0300],\n",
       "          [ 0.9409,  0.4120,  0.6129,  ...,  0.7283, -0.8180,  1.4888],\n",
       "          [ 0.0620,  0.1983,  0.0675,  ..., -0.8164, -2.1203,  0.5834],\n",
       "          ...,\n",
       "          [-0.4766,  0.1509, -1.9746,  ..., -1.1059, -0.4576, -0.5482],\n",
       "          [ 0.1515,  0.2605,  1.2307,  ...,  0.4600,  0.2387, -0.8263],\n",
       "          [-0.2699,  0.1249, -1.4637,  ...,  1.7832, -0.1992, -0.9448]]],\n",
       "\n",
       "\n",
       "        [[[-0.8269, -0.6305,  0.8517,  ...,  0.1559,  0.4775,  0.2800],\n",
       "          [ 0.3817,  1.0025,  0.3259,  ..., -0.0707, -1.1159,  0.0803],\n",
       "          [ 0.3949,  1.8837, -0.5904,  ...,  1.0609, -0.0464,  0.5574],\n",
       "          ...,\n",
       "          [ 1.2029, -0.4394, -0.2655,  ..., -0.4808, -0.0988,  0.9696],\n",
       "          [-1.8645,  0.0060, -0.3333,  ..., -0.6156, -0.9566,  0.2655],\n",
       "          [-1.9624,  0.2632, -0.4373,  ...,  0.5423, -0.6441, -0.5565]]]],\n",
       "       requires_grad=True)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_loss, torch_loss = my_maxpool.mean(), torch_maxpool.mean()\n",
    "my_loss.backward()\n",
    "\n",
    "torch_dL_dX, torch_dL_dO = x.grad, my_maxpool.grad.clone()\n",
    "my_dL_dX = maxpool_backward(torch_dL_dO, x_shape, indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_forward(\n",
    "    x:Tensor,      # (B, fi)\n",
    "    wie:Tensor,    # (fo, fi)\n",
    "    bias:Tensor,   # (fo,)\n",
    "):\n",
    "    return x @ wie.T + bias.unsqueeze(0) \n",
    "\n",
    "def linear_backward(\n",
    "    x:Tensor,       # (B, fi)\n",
    "    wie:Tensor,     # (fo, fi)\n",
    "    bias:Tensor,    # (fo,)\n",
    "    dL_dO:Tensor,   # (B, fo)\n",
    "):\n",
    "    dL_dx = dL_dO @ wie      # (B, fi) <= (B, fo) @ (fo, fi)\n",
    "    dL_dwie = dL_dO.T @ x    # (fo, fi) <= (B, fo).T @ (B, fi)\n",
    "    dL_db = dL_dO.sum(dim=0) # (fo,) <= (B, fo)\n",
    "    return dL_dx, dL_dwie, dL_db"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reshape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reshape_forward(x:Tensor, shape:tuple): # (B, C, H, W)\n",
    "    return x.reshape(shape), x.shape        # (B, C*H*W)\n",
    "\n",
    "def reshape_backward(dL_dO:Tensor, x_shape:torch.Size): # (B, C*H*W)\n",
    "    return dL_dO.reshape(x_shape)                       # (B, C, H, W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ReLU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu_forward(x:Tensor):\n",
    "    return torch.maximum(x, torch.tensor(0))\n",
    "\n",
    "def relu_backward(relu:Tensor, dL_dO:Tensor):\n",
    "    dO_dx = relu * dL_dO\n",
    "    dL_dx = dL_dO * dO_dx\n",
    "    return dL_dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SoftMax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax_forward(logits:Tensor):\n",
    "    max_val, idx = logits.max(-1, keepdim=True) \n",
    "    logits -= max_val\n",
    "    exp = torch.exp(logits)\n",
    "    proba = exp / exp.sum(-1, keepdim=True)\n",
    "    return proba\n",
    "\n",
    "def softmax_backward(probs:Tensor, dL_dprobs:Tensor):\n",
    "    nc = probs.shape[-1]\n",
    "    t1 = torch.einsum(\"ij,ik->ijk\", probs, probs) # (B, nc, nc)\n",
    "    t2 = torch.einsum(\"ij,jk->ijk\", probs, torch.eye(nc, nc)) # (B, nc, nc)\n",
    "    dprobs_dlogits = t2 - t1 # (B, nc, nc)\n",
    "\n",
    "    dL_dlogits = (dL_dprobs[:, None, :] @ dprobs_dlogits)[:, 0, :] # ((B, 1, nc) @ (B, nc, nc))[:, 0, :]\n",
    "    return dL_dlogits # (B, nc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross Entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cross_entropy_forward(y_true:Tensor, y_proba:Tensor):\n",
    "    log_probs = torch.log(y_proba) # (B, nc)\n",
    "    loss = -log_probs[torch.arange(len(y_true)), y_true].mean()\n",
    "    return loss\n",
    "\n",
    "def cross_entropy_backward(y_true:Tensor, y_proba:Tensor):\n",
    "    B = len(y_true)\n",
    "    dL_dlogprobas = torch.zeros_like(y_proba) # (B, nc)\n",
    "    dL_dlogprobas[torch.arange(B), y_true] = -1/B\n",
    "\n",
    "    dlogprobas_dprobas = 1/y_proba # (B, nc)\n",
    "\n",
    "    dL_dprobas = dL_dlogprobas * dlogprobas_dprobas # (B, nc)\n",
    "    return dL_dprobas # (B, nc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LeNet Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Module:\n",
    "    def forward(self, *args):\n",
    "        raise NotImplementedError\n",
    "    def backward(self, *args):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LeNet(Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.parameters = init_weights(\n",
    "            (6, 1, (5, 5)),\n",
    "            (16, 6, (5, 5)),\n",
    "            (120, 16, (5, 5)),\n",
    "            (84, 120),\n",
    "            (10, 84)\n",
    "        )\n",
    "        (\n",
    "            self.W1, self.B1, # Convolutional Layer 1\n",
    "            self.W2, self.B2, # Convolutional Layer 2\n",
    "            self.W3, self.B3, # Convolutional Layer 3\n",
    "            self.W4, self.B4, # Linear Layer 1\n",
    "            self.W5, self.B5  # Linear Layer 2\n",
    "        ) = self.parameters\n",
    "\n",
    "    def forward(self, x:Tensor): # (B, 1, 32, 32)\n",
    "        self.x = x\n",
    "\n",
    "        z1 = conv2d_forward(self.x, self.W1.wei, self.B1.wei) # (B, 6, 28, 28)\n",
    "        self.h1 = relu_forward(z1) # (B, 6, 28, 28)\n",
    "\n",
    "        self.hp1, (self.ridx1, self.cidx1), self.h1_shape = maxpool2d_forward(self.h1, (2, 2), (2, 2)) # (B, 6, 14, 14)\n",
    "\n",
    "        z2 = conv2d_forward(self.hp1, self.W2.wei, self.B2.wei) # (B, 16, 10, 10) ##################\n",
    "        self.h2 = relu_forward(z2) # (B, 16, 10, 10)\n",
    "\n",
    "        self.hp2, (self.ridx2, self.cidx2), self.h2_shape = maxpool2d_forward(self.h2, (2, 2), (2, 2)) # (B, 16, 5, 5) \n",
    "\n",
    "        self.z3 = conv2d_forward(self.hp2, self.W3.wei, self.B3.wei) # (B, 120, 1, 1) \n",
    "        self.h3 = relu_forward(self.z3) # (B, 120, 1, 1) \n",
    "\n",
    "        h3_reshaped, self.h3_shape = reshape_forward(self.h3, (-1, self.h3.size(1))) # (B, 120, 1, 1) -> (B, 120) \n",
    "\n",
    "        self.z4 = linear_forward(h3_reshaped, self.W4.wei, self.B4.wei) # (B, 84)\n",
    "        self.h4 = relu_forward(self.z4) # (B, 84)\n",
    "\n",
    "        self.z5 = linear_forward(self.h4, self.W5.wei, self.B5.wei) # (B, 10)\n",
    "        self.y_proba = softmax_forward(self.z5) # (B, 10)\n",
    "\n",
    "        return self.y_proba\n",
    "    \n",
    "    def backward(self, y_true:Tensor):\n",
    "        dL_dprobs = cross_entropy_backward(y_true, self.y_proba) # (B, 10)\n",
    "\n",
    "        dL_dz5 = softmax_backward(self.y_proba, dL_dprobs) # (B, 10)\n",
    "        dL_dh4, dL_dW5, dL_dB5 = linear_backward(self.z5, self.W5.wei, self.B5.wei, dL_dz5) # (B, 84)\n",
    "        self.W5.grad, self.B5.grad = dL_dW5, dL_dB5\n",
    "\n",
    "        dL_dz4 = relu_backward(self.h4, dL_dh4) # (B, 84)\n",
    "        dL_h3reshaped, dL_dW4, dL_dB4  = linear_backward(self.z4, self.W4.wei, self.B4.wei, dL_dz4) # (B, 120)\n",
    "        self.W4.grad, self.B4.grad = dL_dW4, dL_dB4\n",
    "\n",
    "        dL_dh3 = reshape_backward(dL_h3reshaped, self.h3_shape) # (B, 120, 1, 1)\n",
    "\n",
    "        dL_dz3 = relu_backward(self.h3, dL_dh3) # (B, 120, 1, 1)\n",
    "        dL_dhp2, dL_dW3, dL_dB3 = conv2d_backward(self.hp2, self.W3.wei, self.B3.wei, dL_dz3) # (B, 16, 5, 5)\n",
    "        self.W3.grad, self.B3.grad = dL_dW3, dL_dB3\n",
    "\n",
    "        dL_dh2 = maxpool_backward(dL_dhp2, self.h2_shape, (self.ridx2, self.cidx2)) # (B, 16, 10, 10)\n",
    "\n",
    "        dL_dz2 = relu_backward(self.h2, dL_dh2) # (B, 16, 10, 10)\n",
    "        dL_dhp1, dL_dW2, dL_dB2 = conv2d_backward(self.hp1, self.W2.wei, self.B2.wei, dL_dz2) # (B, 6, 14, 14)\n",
    "        self.W2.grad, self.B2.grad = dL_dW2, dL_dB2\n",
    "\n",
    "        dL_dh1 = maxpool_backward(dL_dhp1, self.h1_shape, (self.ridx1, self.cidx1)) # (B, 6, 28, 28)\n",
    "\n",
    "        dL_dz1 = relu_backward(self.h1, dL_dh1) # (B, 6, 28, 28)\n",
    "        _, dL_dW1, dL_dB1 = conv2d_backward(self.x, self.W1.wei, self.B1.wei, dL_dz1) # (B, 1, 32, 32)\n",
    "        self.W1.grad, self.B1.grad = dL_dW1, dL_dB1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SGD:\n",
    "    def __init__(self, parameters:list[Wei], lr:float):\n",
    "        self.parameters = parameters\n",
    "        self.lr = lr\n",
    "\n",
    "    def step(self):\n",
    "        for p in self.parameters:\n",
    "            p.wei -= self.lr*p.grad\n",
    "\n",
    "    def zero_grad(self):\n",
    "        for p in self.parameters:\n",
    "            p.grad = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LeNet()\n",
    "optimizer = SGD(model.parameters, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @torch.compile # comment while debugging\n",
    "def one_train_step(X_batch:Tensor, y_batch:Tensor):\n",
    "    y_proba = model.forward(X_batch)\n",
    "    loss = cross_entropy_forward(y_batch, y_proba)\n",
    "    model.backward(y_batch)\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "def _conv2d(\n",
    "    x:Tensor, # (H, W)\n",
    "    w:Tensor, # (h, w)\n",
    "    full:bool=False,\n",
    "    convolve:bool=False\n",
    "):\n",
    "    x = x[None, None, ...] # (1, 1, H, W)\n",
    "    w = w[None, None, ...] # (1, 1, h, w)\n",
    "    if full:\n",
    "        pad_h = w.size(-2) - 1\n",
    "        pad_w = w.size(-1) - 1\n",
    "        x = F.pad(x, (pad_w, pad_w, pad_h, pad_h), mode='constant')\n",
    "    if convolve:\n",
    "        w = w.flip((2, 3))\n",
    "    return F.conv2d(\n",
    "        input=x, # (1, 1, H, W)\n",
    "        weight=w, # (1, 1, h, w)\n",
    "        stride=1\n",
    "    )[0, 0, :, :]\n",
    "\n",
    "def conv2d(\n",
    "    inputs:Tensor, # (H, W)\n",
    "    kernel:Tensor, # (hk, wk)\n",
    "    stride:tuple[int, int] # (sh, sw)\n",
    "):\n",
    "    H, W = inputs.size()\n",
    "    hk, wk = kernel.size()\n",
    "    out = torch.zeros(((H-hk+1)//stride[0], (W-wk+1)//stride[1]))\n",
    "    for i in range(0, H, stride[0]):\n",
    "        for j in range(0, W, stride[1]):\n",
    "            portion = inputs[i:i+hk, j:j+wk]\n",
    "            out[i//stride[0], j//stride[1]] = (portion * kernel).sum()\n",
    "    return out\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}