nbs/02b_cleanup.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "841a013a-2dc4-41b2-9b2e-0f68b52ee208",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -Uqq more_itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a080d670-7909-456b-9cf9-5d6813951f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "class InitMixin:\n",
    "    @classmethod\n",
    "    def fmt_init(cls): \n",
    "        sig = inspect.signature(cls.__init__)\n",
    "        params = []\n",
    "        for param in sig.parameters.values():\n",
    "            if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD):\n",
    "                continue\n",
    "            if param.name == \"self\":\n",
    "                continue\n",
    "            fmt = param.name\n",
    "            if param.default is not param.empty:\n",
    "                 fmt += f\"={param.default}\"\n",
    "            params.append(fmt)\n",
    "        return \", \".join(params)\n",
    "    def __repr__(self):\n",
    "        return f\"{self.__class__.__name__}({self.fmt_init()})\"\n",
    "    def __str__(self):\n",
    "        return repr(self)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0df7b80d-a8e4-49cf-9b86-0ad33f9f970f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Module(InitMixin):\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        params = list(inspect.signature(self.forward).parameters.values())\n",
    "        output = self.forward(*args, **kwargs)\n",
    "        inputs = dict(zip([p.name for p in params], list(args) + list(kwargs.values())))\n",
    "        self.cache(**inputs, out=output)\n",
    "        return output\n",
    "    def cache(self, **kwargs):\n",
    "        self.__dict__.update(**kwargs)\n",
    "    def get(self, key, *keys):\n",
    "        keys = [key] + list(keys)\n",
    "        return [self.__dict__.get(k) for k in keys]\n",
    "    def backward(self):\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6a03f2-84a7-4147-9b2d-34106eb2f38b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class New(Module):\n",
    "    def __init__(self, a, b, *args, c=3, **kwargs):\n",
    "        self.a, self.b, self.c = a, b, c\n",
    "    def forward(self, x, y, z=1):\n",
    "        return x + y*z\n",
    "    \n",
    "new = New(1, 2)\n",
    "print(new(1, 2, 3))\n",
    "print(new.__dict__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "639c049e-91f9-4089-98f3-600ddacd2afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import tensor as t\n",
    "from torch.nn import functional as F\n",
    "\n",
    "def test(t1, t2, eps=1e-8): assert torch.allclose(t1, t2, atol=eps)\n",
    "\n",
    "def grad(t): return t.requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76560e6f-2284-4c1b-9f7a-5af16d5615e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1e5770-27e1-4f0a-bb85-d30bbab1c976",
   "metadata": {},
   "source": [
    "## MSE Loss\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "L &= MSE(y, \\hat{y}) = \\mathbb{E}(y - \\hat{y})^2 = \\mathbb{E}(d)^2 \\\\\n",
    "\\frac{dL}{d\\hat{y}} &= 2 \\times \\mathbb{E}(d) = 2\\frac{d}{N}\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9a3878e-59b0-4017-9495-fdfe51459ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MSE(Module):\n",
    "    def forward(self, pred, gt):\n",
    "        return (pred - gt).pow(2).mean()\n",
    "    def backward(self):\n",
    "        d = self.pred - self.gt\n",
    "        self.pred.g = 2 * d / d.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18db3e6f-3666-4979-ba71-f68721d37376",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = grad(t([1.5, 0.3, 2.0])), grad(t([1., 1., 3.]))\n",
    "\n",
    "mse = MSE()\n",
    "my, ref = mse(x, y), F.mse_loss(x, y)\n",
    "test(my, ref)\n",
    "\n",
    "mse.backward()\n",
    "ref.backward()\n",
    "test(x.g, x.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9527b3cf-902b-440f-a013-4d1dbf4f0b9e",
   "metadata": {},
   "source": [
    "## ReLU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e180ae-b442-4c6f-83a7-cb859436e0b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReLU(Module):\n",
    "    def forward(self, inp): return inp.clamp_min(0)\n",
    "    def backward(self): self.inp.g = (self.inp > 0).float() * self.out.g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12400fbb-745c-409d-8549-82d1bb0c1e78",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = grad(t([-1.5, 0., 2.0])), grad(t([1., 1., 3.]))\n",
    "\n",
    "F.mse_loss(F.relu(x), y).backward()\n",
    "\n",
    "mse, relu = MSE(), ReLU()\n",
    "mse(relu(x), y)\n",
    "mse.backward()\n",
    "relu.backward()\n",
    "\n",
    "test(x.g, x.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab922fb-7d8c-4522-9225-c5bbb97a5000",
   "metadata": {},
   "source": [
    "## Linear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62db1094-4650-4179-9fcb-28f28978ec65",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Linear(Module):\n",
    "    def __init__(self, W, b):\n",
    "        super().__init__()\n",
    "        self.W, self.b = W, b\n",
    "    def forward(self, inp): return inp @ self.W + self.b\n",
    "    def backward(self):\n",
    "        i, o = self.inp, self.out\n",
    "        i.g = o.g @ self.W.t()\n",
    "        self.W.g = i.t() @ o.g\n",
    "        self.b.g = o.g.sum(0)  # same bias for all rows\n",
    "    def __repr__(self):\n",
    "        return f\"{self.__class__.__name__}({list(self.W.shape)})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341184ee-e15f-4fb1-a71a-8f8af922d638",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.tensor([\n",
    "    [ 1, -2, 3, 1],\n",
    "    [ 0,  1, 0, 1],\n",
    "    [-1,  0, 1, 0]\n",
    "]).float()\n",
    "y = torch.tensor([1, 0, 3]).unsqueeze(-1)\n",
    "X, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "142e7c8e-d951-4098-8cc8-f2ffdeca84b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "W1 = grad(torch.eye(4))\n",
    "b1 = grad(torch.ones(X.shape[0], W1.shape[-1]))\n",
    "W2 = grad(t([1., 0., -1., 1.]).unsqueeze(-1))\n",
    "b2 = grad(torch.ones((X.shape[0], W2.shape[1])))\n",
    "\n",
    "(F.relu(F.relu(X@W1 + b1)@W2 + b2) - y).pow(2).mean().backward()\n",
    "\n",
    "mse = MSE()\n",
    "lin1, rel1 = Linear(W1, b1), ReLU()\n",
    "lin2, rel2 = Linear(W2, b2), ReLU()\n",
    "\n",
    "mse(rel2(lin2(rel1(lin1(X)))), y)\n",
    "mse.backward()\n",
    "rel2.backward()\n",
    "lin2.backward()\n",
    "rel1.backward()\n",
    "lin1.backward()\n",
    "\n",
    "[test(x.g, x.grad) for x in (W1, W2)]\n",
    "[test(x.g, x.grad.sum(0)) for x in (b1, b2)];"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52c1eba9-62ed-42bb-a38a-87d5724d6f26",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af10ad73-b090-4a96-994b-9b2ec33277ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import reduce\n",
    "class MLP(Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        self.layers = layers\n",
    "    def forward(self, inp):\n",
    "        return reduce(lambda x, l: l(x), self.layers, inp)\n",
    "    def backward(self):\n",
    "        [l.backward() for l in self.layers[::-1]]\n",
    "    def __repr__(self):\n",
    "        fmt = \"\".join(map(lambda l: f\"  {l}\\n\", self.layers))\n",
    "        return f\"MLP(layers=[\\n{fmt}])\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d93d741-1fff-48fc-b700-c51a471495f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "W1 = grad(torch.eye(4))\n",
    "b1 = grad(torch.ones(X.shape[0], W1.shape[-1]))\n",
    "W2 = grad(t([1., 0., -1., 1.]).unsqueeze(-1))\n",
    "b2 = grad(torch.ones(X.shape[0], W2.shape[1]))\n",
    "\n",
    "mse = MSE()\n",
    "lin1, rel1 = Linear(W1, b1), ReLU()\n",
    "lin2, rel2 = Linear(W2, b2), ReLU()\n",
    "model = MLP([lin1, rel1, lin2, rel2])\n",
    "\n",
    "mse(model(X), y)\n",
    "mse.backward()\n",
    "model.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5bfa70c-ed6f-41a6-8d16-7f184cb8f9b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "W1.g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d57ff6c-775b-4239-bb42-f34bc81a14b1",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc67abc3-11f7-489c-960c-8e7708e91250",
   "metadata": {},
   "outputs": [],
   "source": [
    "from more_itertools import chunked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e570b16c-6a5e-41ed-ae1e-4447f942b39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear(i: int, o: int, init=None) -> Linear:\n",
    "    w = torch.randn((i, o))\n",
    "    b = torch.ones(o)\n",
    "    if init is not None: init(w), init(b)\n",
    "    return Linear(w, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94874182-f5ed-4e39-9be3-fd8d4147cfce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(sizes):\n",
    "    layers = []\n",
    "    for i, o in zip(sizes[:-1], sizes[1:]):\n",
    "        layers += [linear(i, o), ReLU()]\n",
    "    layers = layers[:-1]\n",
    "    return MLP(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e96ec047-03ca-4e81-86ae-6eca7e4a8605",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp([784, 10, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc7b14f1-43ef-47f4-bfd2-bb5e69cd3430",
   "metadata": {},
   "outputs": [],
   "source": [
    "MNIST_URL = \"https://github.com/mnielsen/neural-networks-and-deep-learning/blob/master/data/mnist.pkl.gz?raw=true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28b51e87-2025-4ea5-9de4-11f3e6121567",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import pickle\n",
    "from pathlib import Path\n",
    "from urllib.request import urlretrieve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b764065f-18d2-4929-992c-aacef8d367d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mnist(log=True):\n",
    "    path_data = Path(\"data\")\n",
    "    path_data.mkdir(exist_ok=True)\n",
    "    path_gz = path_data/\"mnist.pkl.gz\"\n",
    "    if not path_gz.exists(): \n",
    "        urlretrieve(MNIST_URL, path_gz)\n",
    "    with gzip.open(path_gz, 'rb') as f: \n",
    "        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')\n",
    "    data = x_train, y_train, x_valid, y_valid\n",
    "    if log:\n",
    "        print(f\"train: x.shape={x_train.shape}, y.shape={y_train.shape}\")\n",
    "        print(f\"valid: x.shape={x_valid.shape}, y.shape={y_valid.shape}\")\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7910c60b-2685-4461-bb71-2e274225a1b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision.all import show_images, set_seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c77b7b-f1f4-495e-96e1-85917282fc52",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c3647d-dbd6-498d-b3b6-16a42f4dc6ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1000\n",
    "x_full, y_full, _, _ = get_mnist(False)\n",
    "mask = y_full == 5\n",
    "x_0 = x_full[mask][:n]\n",
    "x_1 = x_full[~mask][:n]\n",
    "y_0 = np.ones(x_0.shape[0])\n",
    "y_1 = np.zeros_like(y_0)\n",
    "x, y = np.vstack([x_0, x_1]), np.hstack([y_0, y_1])\n",
    "idx = np.arange(x.shape[0])\n",
    "np.random.shuffle(idx)\n",
    "x, y = x[idx], y[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db251bb-1e78-4049-a6cb-ecc785b79eab",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([im.reshape(28, 28) for im in x[:15]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b1fef2c-1ce5-46d5-85f8-5aba2d49f7be",
   "metadata": {},
   "outputs": [],
   "source": [
    "from more_itertools import chunked\n",
    "lr = 0.005\n",
    "n_trn = int(x.shape[0] * 0.8)\n",
    "mse = MSE()\n",
    "model = mlp([784, 10, 1])\n",
    "\n",
    "x_trn, y_trn = x[:n_trn], y[:n_trn]\n",
    "x_val, y_val = x[n_trn:], y[n_trn:]\n",
    "\n",
    "for epoch in range(40):\n",
    "    val, cnt = 0, 0\n",
    "    \n",
    "    idx = np.arange(len(x_trn))\n",
    "    np.random.shuffle(idx)\n",
    "    x_trn, y_trn = x_trn[idx], y_trn[idx]\n",
    "    \n",
    "    for xb, yb in zip(chunked(x_trn, 10), chunked(y_trn, 10)):\n",
    "        xb, yb = np.vstack(xb), np.asarray(yb, dtype=np.float32)\n",
    "        xb, yb = torch.as_tensor(xb), torch.as_tensor(yb).view(-1, 1)\n",
    "        logits = model(xb)\n",
    "        val += mse(logits, yb)\n",
    "        cnt += len(yb)\n",
    "        mse.backward()\n",
    "        model.backward()\n",
    "        for lin in (model.layers[0], model.layers[-1]):\n",
    "            lin.W -= lr * lin.W.g\n",
    "            lin.b -= lr * lin.b.g.sum(0)\n",
    "            lin.W.g.zero_()\n",
    "            lin.b.g.zero_()\n",
    "\n",
    "    print(f\"Loss at epoch {epoch+1:2d}: {val/cnt:3.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ccfe7bb-3923-457f-892e-f48dd90ade9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fives = (model(torch.as_tensor(x_val)).clip(0, 1) >= 0.5).sum(1).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec042215-a1b2-42ad-ad1d-d0771c0e8c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "(fives == y_val).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40410eb-c70c-41ce-9902-3ca221a7adf8",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align}\n",
    "f(X, W, b) &= X \\times W + b = H + b\\\\\n",
    "\\\\\n",
    "X_{m,n} &= \n",
    "\\begin{pmatrix}\n",
    "x_{1,1} & x_{1,2} & \\cdots & x_{1,n} \\\\\n",
    "x_{2,1} & x_{2,2} & \\cdots & x_{2,n} \\\\\n",
    "\\vdots  & \\vdots  & \\ddots & \\vdots  \\\\\n",
    "x_{m,1} & x_{m,2} & \\cdots & x_{m,n}\n",
    "\\end{pmatrix}\\\\\n",
    "\\\\\n",
    "W_{n,k} &=\n",
    "\\begin{pmatrix}\n",
    "w_{1,1} & w_{1,2} & \\cdots & w_{1,k} \\\\\n",
    "w_{2,1} & w_{2,2} & \\cdots & w_{2,k} \\\\\n",
    "\\vdots  & \\vdots  & \\ddots & \\vdots  \\\\\n",
    "w_{n,1} & w_{n,2} & \\cdots & w_{n,k}\n",
    "\\end{pmatrix} \\\\\n",
    "\\\\\n",
    "H_{m,k} &=\n",
    "\\begin{pmatrix}\n",
    "h_{1,1} & \\cdots & h_{1,k} \\\\\n",
    "\\vdots  & \\ddots & \\vdots  \\\\\n",
    "h_{m,1} & \\cdots & h_{m,k} \\\\\n",
    "\\end{pmatrix} \\\\\n",
    "\\\\\n",
    "b &= \\left(b_{m,k}\\right)\n",
    "\\end{align}\n",
    "$$"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}