{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from torch import nn, optim\n",
    "from torch.functional import F\n",
    "# enable intellisense\n",
    "%config IPCompleter.greedy=True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_num_correct(preds, labels):\n",
    "    return preds.argmax(dim=1).eq(labels).sum().item()\n",
    "\n",
    "@torch.no_grad()\n",
    "def get_all_preds(model, loader):\n",
    "    all_preds = torch.tensor([])\n",
    "    for batch in loader:\n",
    "        images, labels = batch\n",
    "        preds = model(images)\n",
    "        all_preds = torch.cat(\n",
    "            (all_preds, preds),\n",
    "            dim=0\n",
    "        )\n",
    "    return all_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = torchvision.datasets.MNIST(\n",
    "    root='./data',\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor()\n",
    "    ])\n",
    ")\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=100)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_set = torchvision.datasets.MNIST(\n",
    "    root='./mnist',\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor()\n",
    "    ])\n",
    ")\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train set size 60000\n",
      "train set labels tensor([5, 0, 4,  ..., 5, 6, 8])\n",
      "train label bin count tensor([5923, 6742, 5958, 6131, 5842, 5421, 5918, 6265, 5851, 5949])\n",
      "first image shape torch.Size([1, 28, 28])\n",
      "first image label 5\n",
      "torch.Size([100, 1, 28, 28]) torch.Size([100])\n",
      "labels tensor([5, 0, 4, 1, 9, 2, 1, 3, 1, 4, 3, 5, 3, 6, 1, 7, 2, 8, 6, 9, 4, 0, 9, 1,\n",
      "        1, 2, 4, 3, 2, 7, 3, 8, 6, 9, 0, 5, 6, 0, 7, 6, 1, 8, 7, 9, 3, 9, 8, 5,\n",
      "        9, 3, 3, 0, 7, 4, 9, 8, 0, 9, 4, 1, 4, 4, 6, 0, 4, 5, 6, 1, 0, 0, 1, 7,\n",
      "        1, 6, 3, 0, 2, 1, 1, 7, 9, 0, 2, 6, 7, 8, 3, 9, 0, 4, 6, 7, 4, 6, 8, 0,\n",
      "        7, 8, 3, 1])\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 340376 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "print('train set size', len(train_set))\n",
    "print('train set labels', train_set.targets)\n",
    "print('train label bin count', train_set.targets.bincount())\n",
    "\n",
    "sample = next(iter(train_set))\n",
    "image, label = sample\n",
    "print('first image shape', image.shape)\n",
    "print('first image label', label)\n",
    "\n",
    "\n",
    "batch = next(iter(train_loader))\n",
    "images, labels = batch\n",
    "\n",
    "print(images.shape, labels.shape)\n",
    "\n",
    "grid = torchvision.utils.make_grid(images, nrow=10)\n",
    "\n",
    "plt.figure(figsize=(15, 15))\n",
    "plt.imshow(np.transpose(grid, (1, 2, 0)))\n",
    "print('labels', labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Network(nn.Module):\n",
    "    def __init__(self, kernel_size=5):\n",
    "        super(Network, self).__init__()\n",
    "        \n",
    "        \n",
    "        \n",
    "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=kernel_size)\n",
    "        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=kernel_size)\n",
    "        \n",
    "        \n",
    "        w = h = \n",
    "        self.fc1 = nn.Linear(in_features=12*w*h, out_features=120)\n",
    "        self.fc2 = nn.Linear(in_features=120, out_features=60)\n",
    "        self.out = nn.Linear(in_features=60, out_features=10)\n",
    "        \n",
    "    def forward(self, t):\n",
    "        \n",
    "        # hidden conv layer 1\n",
    "        t = self.conv1(t)\n",
    "        t = F.relu(t)\n",
    "        t = F.max_pool2d(t, kernel_size=2, stride=2)\n",
    "        \n",
    "        # hidden conv layer 2\n",
    "        t = self.conv2(t)\n",
    "        t = F.relu(t)\n",
    "        t = F.max_pool2d(t, kernel_size=2, stride=2)\n",
    "        \n",
    "        # hidden linear layer 1\n",
    "        t = t.reshape(-1, 12 * 4 * 4)\n",
    "        t = self.fc1(t)\n",
    "        t = F.relu(t)\n",
    "        \n",
    "        # hidden linear layer 2\n",
    "        t = self.fc2(t)\n",
    "        t = F.relu(t)\n",
    "        \n",
    "        # output layer\n",
    "        t = self.out(t)\n",
    "        #t = F.softmax(t, dim=1)\n",
    "        \n",
    "        return t\n",
    "        \n",
    "network = Network() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Network(\n",
      "  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2): Conv2d(6, 12, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (fc1): Linear(in_features=192, out_features=120, bias=True)\n",
      "  (fc2): Linear(in_features=120, out_features=60, bias=True)\n",
      "  (out): Linear(in_features=60, out_features=10, bias=True)\n",
      ")\n",
      "conv1.weight \t\t torch.Size([6, 1, 5, 5])\n",
      "conv1.bias \t\t torch.Size([6])\n",
      "conv2.weight \t\t torch.Size([12, 6, 5, 5])\n",
      "conv2.bias \t\t torch.Size([12])\n",
      "fc1.weight \t\t torch.Size([120, 192])\n",
      "fc1.bias \t\t torch.Size([120])\n",
      "fc2.weight \t\t torch.Size([60, 120])\n",
      "fc2.bias \t\t torch.Size([60])\n",
      "out.weight \t\t torch.Size([10, 60])\n",
      "out.bias \t\t torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "print(network)\n",
    "for name, param in network.named_parameters():\n",
    "    print(name, '\\t\\t', param.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CNN output size formula (Squre)\n",
    "\n",
    "* suppose we have an $n \\times n$ input\n",
    "* suppose we have an $f \\times f$ filter\n",
    "* suppose we have a padding of $p$ and a stride of $s$.\n",
    "\n",
    "The output size $O$ is given by this formula:\n",
    "$$\n",
    "O = \\frac{n - f + 2p }{s} + 1\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "single softmax tensor([[0.1278, 0.1114, 0.0736, 0.1042, 0.1032, 0.0571, 0.1200, 0.0946, 0.0926,\n",
      "         0.1155]], grad_fn=<SoftmaxBackward>)\n",
      "initial batch loss 2.213677167892456\n",
      "num of corrent predition tensor(13)\n",
      "network.conv1.weight.grad after backward propagation torch.Size([6, 1, 5, 5])\n",
      "second batch loss 2.208138942718506\n",
      "num of corrent predition tensor(36)\n"
     ]
    }
   ],
   "source": [
    "# pass through single image\n",
    "pred = network(image.unsqueeze(0))\n",
    "print('single softmax', F.softmax(pred, dim=1))\n",
    "\n",
    "\n",
    "optimizer = optim.Adam(network.parameters(), lr=0.001) # Adam optimizer with learning rate as 0.01\n",
    "\n",
    "\n",
    "# pass through a batch\n",
    "preds = network(images)\n",
    "loss = F.cross_entropy(preds, labels)\n",
    "print('initial batch loss', loss.item())\n",
    "print('num of corrent predition', get_num_correct(preds, labels))\n",
    "loss.backward() # backward propapation to update weight gradian for each layer\n",
    "print('network.conv1.weight.grad after backward propagation', network.conv1.weight.grad.shape)\n",
    "optimizer.step()  # update weight\n",
    "\n",
    "# pass the same batch and we should be seeing loss getting better or worse with respect to learning rate\n",
    "preds = network(images)\n",
    "loss = F.cross_entropy(preds, labels)\n",
    "print('second batch loss', loss.item())\n",
    "print('num of corrent predition', get_num_correct(preds, labels))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 total_correct: 52291 loss: 255.15702729672194\n"
     ]
    }
   ],
   "source": [
    "network = Network()\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=100)\n",
    "optimizer = optim.Adam(network.parameters(), lr=0.001)\n",
    "\n",
    "total_loss = 0\n",
    "total_correct = 0\n",
    "\n",
    "for batch in train_loader:\n",
    "    images, labels = batch\n",
    "    preds = network(images)\n",
    "    loss = F.cross_entropy(preds, labels)\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()  # update weights\n",
    "    \n",
    "    total_loss += loss.item()\n",
    "    total_correct += get_num_correct(preds, labels)\n",
    "    \n",
    "print('epoch:', 0, \"total_correct:\", total_correct, \"loss:\", total_loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8715166666666667"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_correct / len(train_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 total_correct: 53130 loss: 238.32186126615852\n",
      "epoch: 1 total_correct: 57952 loss: 67.93871297477745\n",
      "epoch: 2 total_correct: 58567 loss: 46.840392522513866\n",
      "epoch: 3 total_correct: 58863 loss: 36.94275447865948\n",
      "epoch: 4 total_correct: 59071 loss: 31.037055094144307\n",
      "epoch: 5 total_correct: 59196 loss: 26.273945857101353\n",
      "epoch: 6 total_correct: 59295 loss: 22.549735327760573\n",
      "epoch: 7 total_correct: 59424 loss: 19.510454425144417\n",
      "epoch: 8 total_correct: 59517 loss: 17.026573230614304\n",
      "epoch: 9 total_correct: 59545 loss: 15.307725477119675\n"
     ]
    }
   ],
   "source": [
    "network = Network()\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=100)\n",
    "optimizer = optim.Adam(network.parameters(), lr=0.001)\n",
    "\n",
    "\n",
    "\n",
    "for epoch in range(10):\n",
    "    total_loss = 0\n",
    "    total_correct = 0\n",
    "\n",
    "    for batch in train_loader:\n",
    "        images, labels = batch\n",
    "        preds = network(images)\n",
    "        loss = F.cross_entropy(preds, labels)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()  # update weights\n",
    "\n",
    "        total_loss += loss.item()\n",
    "        total_correct += get_num_correct(preds, labels)\n",
    "\n",
    "    print('epoch:', epoch, \"total_correct:\", total_correct, \"loss:\", total_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9924166666666666"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_correct / len(train_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate confusion matrix\n",
    "with torch.no_grad():\n",
    "    prediction_loader = torch.utils.data.DataLoader(train_set, batch_size=10000)\n",
    "    train_preds = get_all_preds(network, prediction_loader)\n",
    "    \n",
    "stacked = torch.stack(\n",
    "    (train_set.targets, train_preds.argmax(dim=1)),\n",
    "    dim=1\n",
    ")\n",
    "print('stacked', stacked)\n",
    "cmt = torch.zeros(10, 10, dtype=torch.int32)\n",
    "for p in stacked:\n",
    "    j, k = p.tolist()\n",
    "    cmt[j, k] += 1\n",
    "print('cmt', cmt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[8.0574e-07, 7.1548e-06, 4.4809e-05, 2.9940e-06, 4.1864e-09, 1.2199e-06,\n",
      "         1.8611e-12, 9.9994e-01, 9.5287e-07, 4.1662e-06]],\n",
      "       grad_fn=<SoftmaxBackward>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fa8604947f0>"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 4840 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# try it on a test sample\n",
    "total_correct / len(train_set)\n",
    "sample = next(iter(test_set))\n",
    "test_image, test_label = sample\n",
    "\n",
    "pred = network(test_image.unsqueeze(0))\n",
    "print(F.softmax(pred, dim=1))\n",
    "plt.imshow(test_image.squeeze(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "??tb.add_scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "batch_size = 100\n",
    "lr = 0.001\n",
    "\n",
    "network = Network()\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size)\n",
    "optimizer = optim.Adam(network.parameters(), lr=lr)\n",
    "\n",
    "images, labels = next(iter(train_loader))\n",
    "grid = torchvision.utils.make_grid(images)\n",
    "\n",
    "comment = f' batch_size={batch_size} lr={lr}'\n",
    "tb = SummaryWriter(comment=comment)\n",
    "tb.add_image('images', torchvision.utils.make_grid(images))\n",
    "tb.add_graph(network, images)\n",
    "\n",
    "\n",
    "network = network.cuda()\n",
    "for epoch in range(2):\n",
    "    total_loss = 0\n",
    "    total_correct = 0\n",
    "    for batch in train_loader:\n",
    "        images, labels = batch\n",
    "        images = images.cuda()\n",
    "        labels = labels.cuda()\n",
    "        preds = network(images)\n",
    "        loss = F.cross_entropy(preds, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "        total_loss += loss.item() * batch_size\n",
    "        total_correct += get_num_correct(preds, labels)\n",
    "        \n",
    "    tb.add_scalar('Loss', total_loss, epoch)\n",
    "    tb.add_scalar('Number Correct', total_correct, epoch)\n",
    "    tb.add_scalar('Accuracy', total_correct / len(train_set), epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Run(lr=0.01, batch_size=100),\n",
       " Run(lr=0.01, batch_size=1000),\n",
       " Run(lr=0.01, batch_size=10000),\n",
       " Run(lr=0.001, batch_size=100),\n",
       " Run(lr=0.001, batch_size=1000),\n",
       " Run(lr=0.001, batch_size=10000)]"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import namedtuple, OrderedDict\n",
    "from itertools import product\n",
    "\n",
    "params = OrderedDict(\n",
    "    lr=[0.01, 0.001],\n",
    "    batch_size=[100, 1000, 10000]\n",
    ")\n",
    "\n",
    "Run = namedtuple('Run', params.keys())\n",
    "runs = [Run(*v) for v in product(*params.values())]\n",
    "runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create on cpu\n",
    "t = torch.ones(1, 1, 28, 28)\n",
    "network = Network()\n",
    "# move to gpu\n",
    "t = t.cuda()\n",
    "network = network.cuda()\n",
    "\n",
    "# move back to cpu\n",
    "t = t.cpu()\n",
    "network = network.cpu()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
