{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Tutorial - The basics\n",
    "\n",
    "### EE-556: Mathematics of Data: From Theory to Computation. Instructor: Prof. Volkan Cevher\n",
    "\n",
    "This tutorial introduces the basics of using the PyTorch library.\n",
    "\n",
    "## What is PyTorch?\n",
    "PyTorch is a deep learning framework that is widely used in research and industry. It is mainly developed by Facebook's AI Research lab (FAIR), and is known for its flexibility and ease of use. Its dynamic computation graph and automatic differentiation capabilities make it especially well-suited for research in machine learning and deep learning.\n",
    "\n",
    "## Installation\n",
    "The installation of PyTorch is straightforward via a package manager like conda (https://conda.io) or pip. The exact installation command depends on your OS, package manager, python version, and whether you have a GPU with CUDA support. In the PyTorch homepage (https://pytorch.org) you can find the correct command in your case. For example, using pip you an simply run the following command in a Linux terminal to install PyTorch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (2.5.1)\n",
      "Requirement already satisfied: torchvision in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (0.20.1)\n",
      "Collecting matplotlib\n",
      "  Downloading matplotlib-3.9.2-cp39-cp39-macosx_11_0_arm64.whl.metadata (11 kB)\n",
      "Requirement already satisfied: filelock in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (4.12.2)\n",
      "Requirement already satisfied: networkx in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (2024.10.0)\n",
      "Requirement already satisfied: sympy==1.13.1 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torch) (1.13.1)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from sympy==1.13.1->torch) (1.3.0)\n",
      "Requirement already satisfied: numpy in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torchvision) (2.0.2)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from torchvision) (11.0.0)\n",
      "Collecting contourpy>=1.0.1 (from matplotlib)\n",
      "  Downloading contourpy-1.3.0-cp39-cp39-macosx_11_0_arm64.whl.metadata (5.4 kB)\n",
      "Collecting cycler>=0.10 (from matplotlib)\n",
      "  Using cached cycler-0.12.1-py3-none-any.whl.metadata (3.8 kB)\n",
      "Collecting fonttools>=4.22.0 (from matplotlib)\n",
      "  Downloading fonttools-4.54.1-cp39-cp39-macosx_11_0_arm64.whl.metadata (163 kB)\n",
      "Collecting kiwisolver>=1.3.1 (from matplotlib)\n",
      "  Downloading kiwisolver-1.4.7-cp39-cp39-macosx_11_0_arm64.whl.metadata (6.3 kB)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from matplotlib) (24.1)\n",
      "Collecting pyparsing>=2.3.1 (from matplotlib)\n",
      "  Downloading pyparsing-3.2.0-py3-none-any.whl.metadata (5.0 kB)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from matplotlib) (2.9.0.post0)\n",
      "Collecting importlib-resources>=3.2.0 (from matplotlib)\n",
      "  Downloading importlib_resources-6.4.5-py3-none-any.whl.metadata (4.0 kB)\n",
      "Requirement already satisfied: zipp>=3.1.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from importlib-resources>=3.2.0->matplotlib) (3.20.2)\n",
      "Requirement already satisfied: six>=1.5 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/fra/math-of-data-2024/.venv/lib/python3.9/site-packages (from jinja2->torch) (3.0.2)\n",
      "Downloading matplotlib-3.9.2-cp39-cp39-macosx_11_0_arm64.whl (7.8 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0mm\n",
      "\u001b[?25hDownloading contourpy-1.3.0-cp39-cp39-macosx_11_0_arm64.whl (249 kB)\n",
      "Using cached cycler-0.12.1-py3-none-any.whl (8.3 kB)\n",
      "Downloading fonttools-4.54.1-cp39-cp39-macosx_11_0_arm64.whl (2.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m19.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading importlib_resources-6.4.5-py3-none-any.whl (36 kB)\n",
      "Downloading kiwisolver-1.4.7-cp39-cp39-macosx_11_0_arm64.whl (64 kB)\n",
      "Downloading pyparsing-3.2.0-py3-none-any.whl (106 kB)\n",
      "Installing collected packages: pyparsing, kiwisolver, importlib-resources, fonttools, cycler, contourpy, matplotlib\n",
      "Successfully installed contourpy-1.3.0 cycler-0.12.1 fonttools-4.54.1 importlib-resources-6.4.5 kiwisolver-1.4.7 matplotlib-3.9.2 pyparsing-3.2.0\n"
     ]
    }
   ],
   "source": [
    "!pip install torch torchvision matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensor creation and manipulation\n",
    "The basic data structure in PyTorch is the tensor. A tensor is a multi-dimensional array that can be used to represent scalars, vectors, matrices, and higher-dimensional arrays. Tensors are similar to numpy arrays, but they have additional features that make them suitable for deep learning tasks. For example, tensors can be moved to a GPU to accelerate computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.randn(10,3,28,28) # usually the shape for images NCHW by default pytorch convention\n",
    "v = torch.rand(28,28)\n",
    "y1 = x * v # elementwise addition, with v replicated across N,C\n",
    "y2 = x @ v == torch.matmul(x,v) #matrix multiplication\n",
    "x = x.permute(0,2,3,1) # tensorflow layout\n",
    "x.reshape(10,3,784) # flatten spatial dimensions\n",
    "x.flatten().shape[0] == x.reshape(-1).shape[0] == (10*3*28*28) # flatten everything"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograd\n",
    "One of the most important features of PyTorch is its automatic differentiation engine, called autograd. Autograd allows you to compute gradients of functions with respect to specified tensors. This is useful for training neural networks using gradient-based optimization algorithms with backpropagation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(10,3,28,28)\n",
    "v = torch.randn(28,28)\n",
    "x.requires_grad = True # <= enable gradient tracking\n",
    "\n",
    "# forward pass\n",
    "A = torch.matmul(x, v) # matrix multiplication \n",
    "L = A.sum() # reduce to scalar\n",
    "# backward pass\n",
    "L.backward() # compute gradients\n",
    "\n",
    "assert v.grad is None # no tracking == no gradient\n",
    "assert x.grad is not None # gradient with respect to x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module API\n",
    "PyTorch provides a high-level API for building neural networks, the `torch.nn.module` API. A simple example is provided below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bias(torch.nn.Module): \n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        # after this is called, can register parameters\n",
    "        self.bias = torch.nn.Parameter(torch.rand(dim), requires_grad=True) # self.bias will now be found by optimisers etc. and have a gradient\n",
    "    \n",
    "    def forward(self, a): \n",
    "        return a + self.bias\n",
    "\n",
    "a = torch.rand(10,5)\n",
    "mu = Bias(5)\n",
    "output = mu(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Question: what is the size of the output variable?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear layers\n",
    "The `torch.nn.Linear` module is used to define a linear transformation with learnable weights. The input and output sizes are specified when the module is created. The forward method is used to compute the output of the layer. Example:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "# Create a linear layer with input size 10 and output size 5\n",
    "linear = nn.Linear(10, 5, bias=True)\n",
    "\n",
    "# Create a random input tensor of size 10\n",
    "x = torch.randn(10)\n",
    "\n",
    "# Compute the output of the linear layer\n",
    "output = linear(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Mathematically, the above is equivalent to the following:\n",
    "$$\n",
    "\\text{output} = W x + b\n",
    "$$\n",
    "where $W$ is a matrix of size $5 \\times 10$ and $b$ is a vector of size $5$.\n",
    "\n",
    "Multi-layer perceptron (MLP): stack several (≥ 2) linear layers, interleaved with activation functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Activation functions\n",
    "Non-linear functions that are applied element-wise and give the neural network its expressivity. Historically sigmoid $\\sigma(x) = \\frac{1}{1 + e^{-x}}$ was common, but due to optimization issues, the rectified linear unit (ReLU) $f(x) = \\max(0, x)$ is the most common choice nowadays. In PyTorch, these activation functions are available as `torch.nn.Sigmoid` and `torch.nn.ReLU`, respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MLP\n",
    "We can now define a simple multi-layer perceptron (MLP) using the `torch.nn.Module` API. The following code defines a simple MLP with one hidden layer and a ReLU activation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(torch.nn.Module):\n",
    "    def __init__(self, dim, hidden_dim=1, out_dim=1):\n",
    "        super().__init__()\n",
    "        # after this is called, can register parameters\n",
    "        self.layer1 = torch.nn.Linear(dim, hidden_dim)\n",
    "        self.layer2 = torch.nn.Linear(hidden_dim, out_dim)\n",
    "        self.activation = torch.nn.ReLU()\n",
    "\n",
    "    def forward(self, a):\n",
    "        hidden = self.activation(self.layer1(a)) \n",
    "        return self.layer2(hidden)\n",
    "\n",
    "a = torch.rand(10,5)\n",
    "mu = MLP(5)\n",
    "output = mu(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolutional layers\n",
    "Convolutional layers are used in computer vision tasks to extract features from images. The `torch.nn.Conv2d` module is used to define a convolutional layer. The input and output channels, kernel size, and stride are specified when the module is created. The forward method is used to compute the output of the layer. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 16, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Define a convolutional layer\n",
    "conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "# Create a random input tensor\n",
    "x = torch.randn(1, 3, 32, 32)\n",
    "\n",
    "# Compute the output of the layer\n",
    "output = conv(x)\n",
    "print(output.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameters of the `nn.Conv2d` module are:\n",
    "- `in_channels`: the number of input channels (e.g., 3 for RGB images)\n",
    "- `out_channels`: the number of output channels (i.e., the number of filters)\n",
    "- `kernel_size`: the size of the convolutional kernel\n",
    "- `stride`: the stride of the convolution, i.e., the number of pixels by which the kernel is shifted\n",
    "- `padding`: the number of pixels to add to the input image before applying the convolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention layers\n",
    "(Self-)attention is a sequence-to-sequence operation that assigns weights to the input sequence elements based on their relevance to each other. Self-attention operates on input sequence $X$ and projects it into three spaces: query $Q$, key $K$, and value $V$:\n",
    "\n",
    "$$\n",
    "Q = XW^Q, \\quad K = XW^K, \\quad V = XW^V\n",
    "$$\n",
    "where $W^Q$, $W^K$, and $W^V$ are learnable weight matrices. The attention matrix is computed as:\n",
    "\n",
    "$$\n",
    "A = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V\n",
    "$$\n",
    "where $d_k$ is the dimension of the key space. The output of the attention layer is finally:\n",
    "$$\n",
    "Y = A V.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 20, 64])\n"
     ]
    }
   ],
   "source": [
    "class SoftmaxAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SoftmaxAttention, self).__init__()\n",
    "        self.wq = nn.Linear(64, 64)\n",
    "        self.wk = nn.Linear(64, 64)\n",
    "        self.wv = nn.Linear(64, 64)\n",
    "\n",
    "    def forward(self, x):\n",
    "        q = self.wq(x)\n",
    "        k = self.wk(x)\n",
    "        v = self.wv(x)\n",
    "\n",
    "        attention = torch.matmul(q, k.transpose(-2, -1))\n",
    "        attention = torch.nn.functional.softmax(attention, dim=-1)\n",
    "        out = torch.matmul(attention, v)\n",
    "        return out\n",
    "\n",
    "attn = SoftmaxAttention()\n",
    "x = torch.randn(10, 20, 64)\n",
    "output = attn(x)\n",
    "print(output.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is known as non-causal (or bidirectional) attention, as it allows each element to attend to all other elements. In the causal (or autoregressive) attention, the attention matrix is maksed to prevent each element from attending to future elements in the input sequence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss functions\n",
    "Express the task that your model is intended to perform on the data\n",
    "\n",
    "- For regression, a sensible default is the mean squareerror $MSE(a,b) = \\frac{1}{N}\\sum_{i=1}^N(a_i - b_i)^2$\n",
    "- For classification, a sensible default is cross-entropy $H(a,b) = -\\sum_{i=1}^N a_i \\log(b_i)$, with $a,b$ being the predicted and true class probabilities, respectively\n",
    "- In PyTorch, `torch.nn.MSELoss` and `torch.nn.CrossEntropyLoss` respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.4129)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Define the loss function\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define the input and target\n",
    "input = torch.randn(3, 5) # (batch_size, num_classes)\n",
    "target = torch.empty(3, dtype=torch.long).random_(5) # (batch_size)\n",
    "\n",
    "# Compute the loss\n",
    "output = loss_fn(input, target)\n",
    "\n",
    "# Print the loss\n",
    "print(output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizer\n",
    "The optimizer is used to update the parameters of the model during training. The most common optimizer is stochastic gradient descent (SGD), but there are many others available in PyTorch, such as Adam, RMSprop, and Adagrad. The optimizer is responsible for updating the weights of the model using the gradients computed by autograd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NestedModel(torch.nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.bias = torch.nn.Parameter(torch.rand(dim), requires_grad=True)\n",
    "        self.mod1 = Bias(dim)\n",
    "        self.mod2 = Bias(dim)\n",
    "        \n",
    "    def forward(self, a):\n",
    "        out = self.mod1(a + self.bias) \n",
    "        return self.mod2(out)\n",
    "\n",
    "a, b = torch.rand(10,5), torch.rand(10,5)\n",
    "model = NestedModel(5)\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.001) # recursively tracks module parameters\n",
    "output = model(a)\n",
    "loss = torch.mean((output - b)**2)\n",
    "optimizer.zero_grad() # clear old gradients\n",
    "loss.backward() # calculate gradients\n",
    "optimizer.step() # one optimization step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting things together: training loop for MLP classification\n",
    "The following code shows a simple training loop for a multi-layer perceptron (MLP) that is trained on the MNIST dataset. The model consists of two linear layers with ReLU activation functions. The loss function is cross-entropy, and the optimizer is stochastic gradient descent (SGD)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 Loss: 6.449501095326742\n",
      "Epoch 1 Loss: 5.902548930358886\n",
      "Epoch 2 Loss: 5.373903276062012\n",
      "Epoch 3 Loss: 4.873682944997151\n",
      "Epoch 4 Loss: 4.414437736256917\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader \n",
    "from torchvision.datasets import MNIST \n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "\n",
    "network = NestedModel(dim=784)\n",
    "optimizer = torch.optim.SGD(network.parameters(),lr=1e-3)\n",
    "cross_entropy = torch.nn.CrossEntropyLoss()\n",
    "EPOCHS = 5\n",
    "for epoch in range(EPOCHS):\n",
    "    loss_avg = []\n",
    "    for a, b in DataLoader(MNIST(root=\".\", transform=transforms.ToTensor(), download=True), batch_size=32):\n",
    "        output = network(a.view(-1, 784)) # flatten the image\n",
    "        loss = cross_entropy(output, b)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "        loss_avg.append(loss.item())\n",
    "    print(f\"Epoch {epoch} Loss: {np.mean(loss_avg)}\")"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}