{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "47d72f0e-bb8b-40f6-9e4f-bfeb1225b1ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting torch\n",
      "  Using cached torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl.metadata (26 kB)\n",
      "Collecting torchvision\n",
      "  Using cached torchvision-0.17.1-cp311-cp311-manylinux1_x86_64.whl.metadata (6.6 kB)\n",
      "Collecting filelock (from torch)\n",
      "  Using cached filelock-3.13.1-py3-none-any.whl.metadata (2.8 kB)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /opt/conda/lib/python3.11/site-packages (from torch) (4.9.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.11/site-packages (from torch) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.11/site-packages (from torch) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.11/site-packages (from torch) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.11/site-packages (from torch) (2023.12.2)\n",
      "Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch)\n",
      "  Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)\n",
      "Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch)\n",
      "  Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)\n",
      "Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch)\n",
      "  Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)\n",
      "Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch)\n",
      "  Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)\n",
      "Collecting nvidia-cublas-cu12==12.1.3.1 (from torch)\n",
      "  Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)\n",
      "Collecting nvidia-cufft-cu12==11.0.2.54 (from torch)\n",
      "  Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)\n",
      "Collecting nvidia-curand-cu12==10.3.2.106 (from torch)\n",
      "  Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl.metadata (1.5 kB)\n",
      "Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch)\n",
      "  Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)\n",
      "Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch)\n",
      "  Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl.metadata (1.6 kB)\n",
      "Collecting nvidia-nccl-cu12==2.19.3 (from torch)\n",
      "  Using cached nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl.metadata (1.8 kB)\n",
      "Collecting nvidia-nvtx-cu12==12.1.105 (from torch)\n",
      "  Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl.metadata (1.7 kB)\n",
      "Collecting triton==2.2.0 (from torch)\n",
      "  Using cached triton-2.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (1.4 kB)\n",
      "Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch)\n",
      "  Downloading nvidia_nvjitlink_cu12-12.4.99-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.11/site-packages (from torchvision) (1.26.3)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.11/site-packages (from torchvision) (10.2.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.11/site-packages (from jinja2->torch) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.11/site-packages (from sympy->torch) (1.3.0)\n",
      "Using cached torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl (755.6 MB)\n",
      "Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB)\n",
      "Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB)\n",
      "Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB)\n",
      "Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB)\n",
      "Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB)\n",
      "Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB)\n",
      "Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB)\n",
      "Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB)\n",
      "Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB)\n",
      "Using cached nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl (166.0 MB)\n",
      "Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB)\n",
      "Using cached triton-2.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (167.9 MB)\n",
      "Using cached torchvision-0.17.1-cp311-cp311-manylinux1_x86_64.whl (6.9 MB)\n",
      "Using cached filelock-3.13.1-py3-none-any.whl (11 kB)\n",
      "Downloading nvidia_nvjitlink_cu12-12.4.99-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m27.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, filelock, triton, nvidia-cusparse-cu12, nvidia-cudnn-cu12, nvidia-cusolver-cu12, torch, torchvision\n",
      "Successfully installed filelock-3.13.1 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.19.3 nvidia-nvjitlink-cu12-12.4.99 nvidia-nvtx-cu12-12.1.105 torch-2.2.1 torchvision-0.17.1 triton-2.2.0\n"
     ]
    }
   ],
   "source": [
    "#!pip3 install torch==1.2.0+cu92 torchvision==0.4.0+cu92 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "!pip3 install torch torchvision\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da9ca25-5156-4216-8ac7-b2347e60c34e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d11c9b95-ea6b-457e-a3e7-c5c8de58b450",
   "metadata": {},
   "outputs": [],
   "source": [
    "## The usual imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "## for printing image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4aad0f0f-545e-416b-97f7-2678c782215a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2.1+cu121\n"
     ]
    }
   ],
   "source": [
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "54c98b09-a734-4edf-96c0-20b78bed9d1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## parameter denoting the batch size\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "## transformations\n",
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor()])\n",
    "\n",
    "## download and load training dataset\n",
    "trainset = torchvision.datasets.MNIST(root='./data', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True, num_workers=2)\n",
    "\n",
    "## download and load testing dataset\n",
    "testset = torchvision.datasets.MNIST(root='./data', train=False,\n",
    "                                       download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,\n",
    "                                         shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "23880f4e-4790-4055-bbb4-1b0c47755b9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2b57e46c50>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 66072 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image = transforms.ToPILImage(mode='L')(torch.randn(1, 96, 96))\n",
    "plt.imshow(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aadba5ab-c84c-4f80-8c79-d30633343a61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2b55422f50>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 65288 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## dummy transformation\n",
    "dummy_transform = transforms.Compose(\n",
    "    [transforms.RandomRotation(45)])\n",
    "\n",
    "dummy_result = dummy_transform(image)\n",
    "\n",
    "plt.imshow(dummy_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "11cd3a28-2caa-411e-8091-4f24dc600ba5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2b55496790>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 62348 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "## dummy transform \n",
    "dummy2_transform = transforms.Compose(\n",
    "    [transforms.RandomRotation(45), transforms.RandomVerticalFlip()])\n",
    "\n",
    "dummy2_result = dummy2_transform(image)\n",
    "\n",
    "plt.imshow(dummy2_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9ad649e7-7617-4658-aa61-34348836427a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85440 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## functions to show an image\n",
    "def imshow(img):\n",
    "    #img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "\n",
    "## get some random training images\n",
    "dataiter = iter(trainloader)\n",
    "images, labels = next(dataiter)\n",
    "\n",
    "## show images\n",
    "imshow(torchvision.utils.make_grid(images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21260ecc-202e-40a7-ad12-02d49055bdc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image batch dimensions: torch.Size([32, 1, 28, 28])\n",
      "Image label dimensions: torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "for images, labels in trainloader:\n",
    "    print(\"Image batch dimensions:\", images.shape)\n",
    "    print(\"Image label dimensions:\", labels.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "00450688-b5a3-4df7-abf5-1428dd7a3792",
   "metadata": {},
   "outputs": [],
   "source": [
    "## the model\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.d1 = nn.Linear(28 * 28, 128)\n",
    "        self.dropout = nn.Dropout(p=0.2)\n",
    "        self.d2 = nn.Linear(128, 10)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x.flatten(start_dim = 1)\n",
    "        x = self.d1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        logits = self.d2(x)\n",
    "        out = F.softmax(logits, dim=1)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "20abb31a-584c-4efc-850a-308853300b14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch size: torch.Size([32, 1, 28, 28])\n",
      "torch.Size([32, 10])\n"
     ]
    }
   ],
   "source": [
    "## test the model with 1 batch\n",
    "model = MyModel()\n",
    "for images, labels in trainloader:\n",
    "    print(\"batch size:\", images.shape)\n",
    "    out = model(images)\n",
    "    print(out.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "71f3b1ca-6507-48d2-bde5-3d7a437fb108",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.001\n",
    "num_epochs = 5\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = MyModel()\n",
    "model = model.to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f3aedefc-1e05-4952-a390-2d7fae26bd2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## utility function to compute accuracy\n",
    "def get_accuracy(output, target, batch_size):\n",
    "    ''' Obtain accuracy for training round '''\n",
    "    corrects = (torch.max(output, 1)[1].view(target.size()).data == target.data).sum()\n",
    "    accuracy = 100.0 * corrects/batch_size\n",
    "    return accuracy.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "498cbc09-c3bd-4e57-9aab-ab192c49198a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## train the model\n",
    "for epoch in range(num_epochs):\n",
    "    train_running_loss = 0.0\n",
    "    train_acc = 0.0\n",
    "\n",
    "    ## commence training\n",
    "    model = model.train()\n",
    "\n",
    "    ## training step\n",
    "    for i, (images, labels) in enumerate(trainloader):\n",
    "        \n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        ## forward + backprop + loss\n",
    "        predictions = model(images)\n",
    "        loss = criterion(predictions, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "\n",
    "        ## update model params\n",
    "        optimizer.step()\n",
    "\n",
    "        train_running_loss += loss.detach().item()\n",
    "        train_acc += get_accuracy(predictions, labels, BATCH_SIZE)\n",
    "    \n",
    "    model.eval()\n",
    "    print('Epoch: %d | Loss: %.4f | Train Accuracy: %.2f' \\\n",
    "          %(epoch, train_running_loss / i, train_acc/i)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4df7d5-b244-4a68-b06a-bb750c975cf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_acc = 0.0\n",
    "for i, (images, labels) in enumerate(testloader, 0):\n",
    "    images = images.to(device)\n",
    "    labels = labels.to(device)\n",
    "    outputs = model(images)\n",
    "    test_acc += get_accuracy(outputs, labels, BATCH_SIZE)\n",
    "        \n",
    "print('Test Accuracy: %.2f'%( test_acc/i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76669cbb-29a0-402e-88ed-9b0ce913945f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
