lec_4_v4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f2601a36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, pickle, gzip, matplotlib as mp\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "mp.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "658da799",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50000, 784]),\n",
       " torch.Size([50000]),\n",
       " torch.Size([10000, 784]),\n",
       " torch.Size([10000]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = None\n",
    "with gzip.open('data/mnist.pkl.gz', 'rb') as f:\n",
    "    data = pickle.load(f, encoding='latin')\n",
    "    \n",
    "((x_train, y_train), (x_val, y_val), _) = data\n",
    "(x_train, y_train, x_val, y_val) = map(torch.tensor, (x_train, y_train, x_val, y_val))\n",
    "x_train.shape, y_train.shape, x_val.shape, y_val.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7b8405e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "m,n = x_train.shape\n",
    "c = max(y_train) + 1\n",
    "nh = 50\n",
    "nout = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "e4e9e415",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        super().__init__()\n",
    "        self.layers = [nn.Linear(nin, nh), nn.ReLU(nh), nn.Linear(nh, nout)]\n",
    "    \n",
    "    def __call__(self, x):\n",
    "        for l in self.layers:\n",
    "            x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "652418cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(n, nh, nout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "894ac6e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model()"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "6ee71319",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Linear(in_features=784, out_features=50, bias=True),\n",
       " ReLU(inplace=True),\n",
       " Linear(in_features=50, out_features=10, bias=True)]"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a1a9d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1a128f30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50000, 10])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a8d10d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logsoftmax(x):\n",
    "    return x.exp()/(x.exp().sum(1, keepdim=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "d62162e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logsoftmax(x):\n",
    "    return x - (x.exp().sum(1, keepdim=True)).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "6cf68f10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-23.0668, grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logsoftmax(pred)[0].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "9e292b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logsumexp(x):\n",
    "    m = x.max(1, keepdim=True)[0]\n",
    "    return m + ((x-m).exp().sum(1, keepdim=True)).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "d0bc6f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logsoftmax(x):\n",
    "    return x - logsumexp(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "3ae9b724",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-23.0668, grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logsoftmax(pred)[0].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "8a36a16e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(logsumexp(pred), pred.logsumexp(1, keepdim=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "64ba41dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nll(pred, target):\n",
    "    return -pred[range(pred.shape[0]), target].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "a5341da7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0509, grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nll(pred, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "a376dd98",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_pred = logsoftmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "c06627b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.2991, grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nll(sm_pred, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d7f29fa1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.2991, grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = nll(sm_pred, y_train)\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "d3b12cf7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(loss, F.nll_loss(sm_pred, y_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "a31c6204",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(loss, F.cross_entropy(sm_pred, y_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "82536222",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = F.cross_entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "9eb35ab3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 784]), torch.Size([50]))"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 50\n",
    "xb, yb = x_train[:bs], y_train[:bs]\n",
    "xb.shape, yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "c7508aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = model(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "e9cda4ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50, 10])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "f7ab1001",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3032, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = loss_func(preds, yb)\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "099f257b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(preds, target):\n",
    "    return (preds.argmax(dim=1) == target).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "6e2402bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0800)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy(preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "5ab929c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def report(loss, preds, target):\n",
    "    print(f'loss:{loss:.2f}, accuracy:{accuracy(preds, target):.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "37b70315",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:2.30, accuracy:0.08\n"
     ]
    }
   ],
   "source": [
    "report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "fa88412c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.5\n",
    "epochs = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "d3cab229",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:0.13, accuracy:0.96\n",
      "loss:0.08, accuracy:0.98\n",
      "loss:0.07, accuracy:0.98\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(0, m, bs):\n",
    "        s = slice(i, min(m, i+bs))\n",
    "        xb,yb = x_train[s],y_train[s]\n",
    "        preds = model(xb)\n",
    "        loss = loss_func(preds, yb)\n",
    "        loss.backward()\n",
    "        with torch.no_grad():\n",
    "            for l in model.layers:\n",
    "                if hasattr(l, 'weight'):\n",
    "                    l.weight -= lr * l.weight.grad\n",
    "                    l.bias -= lr * l.bias.grad\n",
    "                    l.weight.grad.zero_()\n",
    "                    l.bias.grad.zero_()\n",
    "    report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "b715a22c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        super().__init__()\n",
    "        self.l1 = nn.Linear(nin, nh)\n",
    "        self.relu = nn.ReLU(nh)\n",
    "        self.l2 = nn.Linear(nh, nout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.l2(self.relu(self.l1(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "9b1ce257",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLP(n, nh, nout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "efcdaa69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MLP(\n",
       "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "c67e9ac9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MLP(\n",
       "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "4ddfe243",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for i in range(0, m, bs):\n",
    "            s = slice(i, min(m, i+bs))\n",
    "            xb,yb = x_train[s],y_train[s]\n",
    "            preds = model(xb)\n",
    "            loss = loss_func(preds, yb)\n",
    "            loss.backward()\n",
    "            with torch.no_grad():\n",
    "                for p in model.parameters():\n",
    "                    p -= p.grad * lr\n",
    "                model.zero_grad()\n",
    "        report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "05d26c6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:0.08, accuracy:0.96\n",
      "loss:0.07, accuracy:0.96\n",
      "loss:0.05, accuracy:0.98\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "8d422a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModule:\n",
    "    def __init__(self):\n",
    "        self._modules = {}\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return f'{self._modules}'\n",
    "    \n",
    "    def __setattr__(self, k, v):\n",
    "        if not k.startswith('_'):\n",
    "            self._modules[k] = v\n",
    "        super().__setattr__(k, v)\n",
    "        \n",
    "    def parameters(self):\n",
    "        for l in self._modules.values():\n",
    "            yield from l.parameters()\n",
    "            \n",
    "    def zero_grad(self):\n",
    "        for p in self.parameters():\n",
    "            p.grad.zero_()\n",
    "            \n",
    "    def __call__(self, x):\n",
    "        return self.forward(x)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "d114f374",
   "metadata": {},
   "outputs": [],
   "source": [
    "md1 = MyModule()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "fd0ecf71",
   "metadata": {},
   "outputs": [],
   "source": [
    "md1.foo = nn.Linear(3, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "43314aec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'foo': Linear(in_features=3, out_features=4, bias=True)}"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "md1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "a7848df2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 3])\n",
      "torch.Size([4])\n"
     ]
    }
   ],
   "source": [
    "for p in md1.parameters():\n",
    "    print(p.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "569537cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(MyModule):\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        super().__init__()\n",
    "        self.l1 = nn.Linear(nin, nh)\n",
    "        self.relu = nn.ReLU(nh)\n",
    "        self.l2 = nn.Linear(nh, nout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.l2(self.relu(self.l1(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "6b14c210",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLP(n, nh, nout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "0dfbac0d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'l1': Linear(in_features=784, out_features=50, bias=True), 'relu': ReLU(inplace=True), 'l2': Linear(in_features=50, out_features=10, bias=True)}"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "3f79ce61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:0.17, accuracy:0.92\n",
      "loss:0.10, accuracy:0.98\n",
      "loss:0.10, accuracy:0.96\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92767f65",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:python_main] *",
   "language": "python",
   "name": "conda-env-python_main-py"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}