hw2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 10-714 Homework 2\n",
    "\n",
    "In this homework, you will be implementing a neural network library in the needle framework. Reminder: __you must save a copy in drive__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to set up the assignment\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "%cd /content/drive/MyDrive/\n",
    "!mkdir -p 10714\n",
    "%cd /content/drive/MyDrive/10714\n",
    "!git clone https://github.com/dlsys10714/hw2.git\n",
    "%cd /content/drive/MyDrive/10714/hw2\n",
    "\n",
    "!pip3 install --upgrade --no-deps git+https://github.com/dlsys10714/mugrade.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 0\n",
    "\n",
    "This homework builds off of Homework 1. First, in your Homework 2 directory, copy the files `python/needle/autograd.py`, `python/needle/ops/ops_mathematic.py` from your Homework 1.\n",
    "\n",
    "***NOTE***: The default data type for the tensor is `float32`. If you want to change the data type, you can do so by setting the `dtype` parameter in the `Tensor` constructor. For example, `Tensor([1, 2, 3], dtype='float64')` will create a tensor with `float64` data type. \n",
    "In this homework, **make sure any tensor you create has `float32` data type to avoid any issues with the autograder**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./python')\n",
    "sys.path.append('./apps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Question 1\n",
    "\n",
    "In this first question, you will implement a few different methods for weight initialization.  This will be done in the `python/needle/init/init_initializers.py` file, which contains a number of routines for initializing needle Tensors using various random and constant initializations.  Following the same methodology of the existing initializers (you will want to call e.g. `init.rand` or `init.randn` implemented in `python/needle/init/init_basic.py` from your functions below, implement the following common initialization methods.  In all cases, the functions should return `fan_in` by `fan_out` 2D tensors (extensions to other sizes can be done via e.g., reshaping).\n",
    "\n",
    "\n",
    "### Xavier uniform\n",
    "`xavier_uniform(fan_in, fan_out, gain=1.0, **kwargs)`\n",
    "\n",
    "Fills the input Tensor with values according to the method described in [Understanding the difficulty of training deep feedforward neural networks](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf), using a uniform distribution. The resulting Tensor will have values sampled from $\\mathcal{U}(-a, a)$ where \n",
    "\\begin{equation}\n",
    "a = \\text{gain} \\times \\sqrt{\\frac{6}{\\text{fan_in} + \\text{fan_out}}}\n",
    "\\end{equation}\n",
    "\n",
    "Pass remaining `**kwargs` parameters to the corresponding `init` random call.\n",
    "\n",
    "##### Parameters\n",
    "- `fan_in` - dimensionality of input\n",
    "- `fan_out` - dimensionality of output\n",
    "- `gain` - optional scaling factor\n",
    "___\n",
    "\n",
    "### Xavier normal\n",
    "`xavier_normal(fan_in, fan_out, gain=1.0, **kwargs)`\n",
    "\n",
    "Fills the input Tensor with values according to the method described in [Understanding the difficulty of training deep feedforward neural networks](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf), using a normal distribution. The resulting Tensor will have values sampled from $\\mathcal{N}(0, \\text{std}^2)$ where \n",
    "\\begin{equation}\n",
    "\\text{std} = \\text{gain} \\times \\sqrt{\\frac{2}{\\text{fan_in} + \\text{fan_out}}}\n",
    "\\end{equation}\n",
    "\n",
    "##### Parameters\n",
    "- `fan_in` - dimensionality of input\n",
    "- `fan_out` - dimensionality of output\n",
    "- `gain` - optional scaling factor\n",
    "___\n",
    "\n",
    "### Kaiming uniform\n",
    "`kaiming_uniform(fan_in, fan_out, nonlinearity=\"relu\", **kwargs)`\n",
    "\n",
    "Fills the input Tensor with values according to the method described in [Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification](https://arxiv.org/pdf/1502.01852.pdf), using a uniform distribution. The resulting Tensor will have values sampled from $\\mathcal{U}(-\\text{bound}, \\text{bound})$ where \n",
    "\\begin{equation}\n",
    "\\text{bound} = \\text{gain} \\times \\sqrt{\\frac{3}{\\text{fan_in}}}\n",
    "\\end{equation}\n",
    "\n",
    "Use the recommended gain value for ReLU: $\\text{gain}=\\sqrt{2}$.\n",
    "\n",
    "##### Parameters\n",
    "- `fan_in` - dimensionality of input\n",
    "- `fan_out` - dimensionality of output\n",
    "- `nonlinearity` - the non-linear function\n",
    "___\n",
    "\n",
    "### Kaiming normal\n",
    "`kaiming_normal(fan_in, fan_out, nonlinearity=\"relu\", **kwargs)`\n",
    "\n",
    "Fills the input Tensor with values according to the method described in [Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification](https://arxiv.org/pdf/1502.01852.pdf), using a uniform distribution. The resulting Tensor will have values sampled from $\\mathcal{N}(0, \\text{std}^2)$ where \n",
    "\\begin{equation}\n",
    "\\text{std} = \\frac{\\text{gain}}{\\sqrt{\\text{fan_in}}}\n",
    "\\end{equation}\n",
    "\n",
    "Use the recommended gain value for ReLU: $\\text{gain}=\\sqrt{2}$.\n",
    "\n",
    "##### Parameters\n",
    "- `fan_in` - dimensionality of input\n",
    "- `fan_out` - dimensionality of output\n",
    "- `nonlinearity` - the non-linear function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_init\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"init\" -s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 2\n",
    "\n",
    "In this question, you will implement additional modules in `python/needle/nn/nn_basic.py`. Specifically, for the following modules described below, initialize any variables of the module in the constructor, and fill out the `forward` method. **Note:** Be sure that you are using the `init` functions that you just implemented to initialize the parameters, and don't forget to pass the `dtype` argument.\n",
    "___\n",
    "\n",
    "### Linear\n",
    "`needle.nn.Linear(in_features, out_features, bias=True, device=None, dtype=\"float32\")`\n",
    "\n",
    "Applies a linear transformation to the incoming data: $y = xA^T + b$. The input shape is $(N, H_{in})$ where $H_{in}=\\text{in_features}$. The output shape is $(N, H_{out})$ where $H_{out}=\\text{out_features}$.\n",
    "\n",
    "**Be careful to explicitly broadcast the bias term to the correct shape -- Needle does not support implicit broadcasting.**\n",
    "\n",
    "**Note: for all layers including this one, you should initialize the weight Tensor before the bias Tensor, and should initialize all Parameters using only functions from `init`**. This does not affect the algorithm's correctness. It is only necessary to ensure the value matches the expected results in the mugrade tests for this assignment's implementation scope. \n",
    "\n",
    "##### Parameters\n",
    "- `in_features` - size of each input sample\n",
    "- `out_features` - size of each output sample\n",
    "- `bias` - If set to `False`, the layer will not learn an additive bias.\n",
    "\n",
    "##### Variables\n",
    "- `weight` - the learnable weights of shape (`in_features`, `out_features`). The values should be initialized with the Kaiming Uniform initialization with `fan_in = in_features`\n",
    "- `bias` - the learnable bias of shape (`out_features`). The values should be initialized with the Kaiming Uniform initialize with `fan_in = out_features`. **Note the difference in fan_in choice, due to their relative sizes**. \n",
    "\n",
    "Make sure to enclose all necessary variables e.g. (`weight`, `bias`) in the `Parameter` class so that they are visible to the optimizers which would be implemented next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_linear\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_linear\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ReLU\n",
    "`needle.nn.ReLU()`\n",
    "\n",
    "Applies the rectified linear unit function element-wise:\n",
    "$ReLU(x) = max(0, x)$.\n",
    "\n",
    "If you have previously implemented ReLU's backwards pass in terms of itself, note that this is numerically unstable and will likely cause problems\n",
    "down the line.\n",
    "Instead, consider that we could write the derivative of ReLU as $I\\{x>0\\}$, where we arbitrarily decide that the derivative at $x=0$ is 0.\n",
    "(This is a _subdifferentiable_ function.)\n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_relu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_relu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Sequential\n",
    "`needle.nn.Sequential(*modules)`\n",
    "\n",
    "Applies a sequence of modules to the input (in the order that they were passed to the constructor) and returns the output of the last module.\n",
    "These should be kept in a `.module` property: you should _not_ redefine any magic methods like `__getitem__`, as this may not be compatible with our tests.\n",
    "\n",
    "##### Parameters\n",
    "- `*modules` - any number of modules of type `needle.nn.Module`\n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_sequential\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_sequential\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### LogSumExp\n",
    "\n",
    "`needle.ops.LogSumExp(axes)`\n",
    "\n",
    "Applies a numerically stable log-sum-exp function to the input by subtracting off the maximum elements. You will need to implement this and the next operation in file `python/needle/ops/ops_logarithmic.py`.\n",
    "\n",
    "\\begin{equation}\n",
    "\\text{LogSumExp}(z) = \\log (\\sum_{i} \\exp (z_i - \\max{z})) + \\max{z}\n",
    "\\end{equation}\n",
    "\n",
    "#### Parameters\n",
    "- `axes` - Tuple of axes to sum and take the maximum element over. This uses the same conventions as `needle.ops.Summation()`\n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_op_logsumexp\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"op_logsumexp\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LogSoftmax\n",
    "\n",
    "`needle.ops.LogSoftmax(axes)`\n",
    "\n",
    "Applies a numerically stable logsoftmax function to the input by subtracting off the maximum elements. Assume the input NDArray is 2 dimensional and we are doing softmax over `axis=1`.\n",
    "\n",
    "\\begin{equation}\n",
    "\\text{LogSoftmax}(z) = \\log \\left(\\frac{\\exp(z_i - \\max z)}{\\sum_{i}\\exp(z_i - \\max z)}\\right) = z - \\text{LogSumExp}(z)\n",
    "\\end{equation}\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_op_logsoftmax\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"op_logsoftmax\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### SoftmaxLoss\n",
    "\n",
    "`needle.nn.SoftmaxLoss()`\n",
    "\n",
    "Applies the softmax loss as defined below (and as implemented in Homework 1), taking in as input a Tensor of logits and a Tensor of the true labels (expressed as a list of numbers, *not* one-hot encoded).\n",
    "\n",
    "Note that you can use the `init.one_hot` function now instead of writing this yourself.  Note: You will need to use the numerically stable logsumexp operator you just implemented for this purpose.\n",
    "\n",
    "\\begin{equation}\n",
    "\\ell_\\text{softmax}(z,y) = \\log \\sum_{i=1}^k \\exp z_i - z_y\n",
    "\\end{equation}\n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_softmax_loss\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_softmax_loss\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### LayerNorm1d\n",
    "`needle.nn.LayerNorm1d(dim, eps=1e-5, device=None, dtype=\"float32\")`\n",
    "\n",
    "Applies layer normalization over a mini-batch of inputs as described in the paper [Layer Normalization](https://arxiv.org/abs/1607.06450).\n",
    "\n",
    "\\begin{equation}\n",
    "y = w \\circ \\frac{x_i - \\textbf{E}[x]}{((\\textbf{Var}[x]+\\epsilon)^{1/2})} + b\n",
    "\\end{equation}\n",
    "\n",
    "where $\\textbf{E}[x]$ denotes the empirical mean of the inputs, $\\textbf{Var}[x]$ denotes their empirical variance (note that here we are using the \"biased\" estimate of the variance, i.e., dividing by $N$ rather than by $N-1$), and $w$ and $b$ denote learnable scalar weights and biases respectively.  Note you can assume the input to this layer is a 2D tensor, with batches in the first dimension and features in the second. You might need to broadcast the weight and bias before applying them.\n",
    "\n",
    "##### Parameters\n",
    "- `dim` - number of channels\n",
    "- `eps` - a value added to the denominator for numerical stability.\n",
    "\n",
    "##### Variables\n",
    "- `weight` - the learnable weights of size `dim`, elements initialized to 1.\n",
    "- `bias` - the learnable bias of shape `dim`, elements initialized to 0.\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_layernorm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_layernorm\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Flatten\n",
    "`needle.nn.Flatten()`\n",
    "\n",
    "Takes in a tensor of shape `(B,X_0,X_1,...)`, and flattens all non-batch dimensions so that the output is of shape `(B, X_0 * X_1 * ...)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_flatten\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_flatten\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BatchNorm1d\n",
    "`needle.nn.BatchNorm1d(dim, eps=1e-5, momentum=0.1, device=None, dtype=\"float32\")`\n",
    "\n",
    "Applies batch normalization over a mini-batch of inputs as described in the paper [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167).\n",
    "\n",
    "\\begin{equation}\n",
    "y = w \\circ \\frac{z_i - \\textbf{E}[x]}{((\\textbf{Var}[x]+\\epsilon)^{1/2})} + b\n",
    "\\end{equation}\n",
    "\n",
    "but where here the mean and variance refer to to the mean and variance over the _batch_dimensions.  The function also computes a running average of mean/variance for all features at each layer $\\hat{\\mu}, \\hat{\\sigma}^2$, and at test time normalizes by these quantities:\n",
    "\n",
    "\\begin{equation}\n",
    "y = \\frac{(x - \\hat{mu})}{((\\hat{\\sigma}^2_{i+1})_j+\\epsilon)^{1/2}}\n",
    "\\end{equation}\n",
    "\n",
    "\n",
    "BatchNorm uses the running estimates of mean and variance instead of batch statistics at test time, i.e.,\n",
    "after `model.eval()` has been called on the BatchNorm layer's `training` flag is false.\n",
    "\n",
    "To compute the running estimates, you can use the equation $$\\hat{x_{new}} = (1 - m) \\hat{x_{old}} + mx_{observed},$$\n",
    "where $m$ is momentum.\n",
    "\n",
    "##### Parameters\n",
    "- `dim` - input dimension\n",
    "- `eps` - a value added to the denominator for numerical stability.\n",
    "- `momentum` - the value used for the running mean and running variance computation.\n",
    "\n",
    "##### Variables\n",
    "- `weight` - the learnable weights of size `dim`, elements initialized to 1.\n",
    "- `bias` - the learnable bias of size `dim`, elements initialized to 0.\n",
    "- `running_mean` - the running mean used at evaluation time, elements initialized to 0.\n",
    "- `running_var` - the running (unbiased) variance used at evaluation time, elements initialized to 1. \n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_batchnorm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_batchnorm\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dropout\n",
    "`needle.nn.Dropout(p = 0.5)`\n",
    "\n",
    "During training, randomly zeroes some of the elements of the input tensor with probability `p` using samples from a Bernoulli distribution. This has proven to be an effective technique for regularization and preventing the co-adaptation of neurons as described in the paper [Improving neural networks by preventing co-adaption of feature detectors](https://arxiv.org/abs/1207.0580). During evaluation the module simply computes an identity function. \n",
    "\n",
    "\\begin{equation}\n",
    "\\hat{z}_{i+1} = \\sigma_i (W_i^T z_i + b_i) \\\\\n",
    "(z_{i+1})_j = \n",
    "    \\begin{cases}\n",
    "    (\\hat{z}_{i+1})_j /(1-p) & \\text{with probability } 1-p \\\\\n",
    "    0 & \\text{with probability } p \\\\\n",
    "    \\end{cases}\n",
    "\\end{equation}\n",
    "\n",
    "**Important**: If the Dropout module the flag `training=False`, you shouldn't \"dropout\" any weights. That is, dropout applies during training only, not during evaluation. Note that `training` is a flag in `nn.Module`.\n",
    "\n",
    "##### Parameters\n",
    "- `p` - the probability of an element to be zeroed.\n",
    "\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_dropout\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_dropout\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Residual\n",
    "`needle.nn.Residual(fn: Module)`\n",
    "\n",
    "Applies a residual or skip connection given module $\\mathcal{F}$ and input Tensor $x$, returning $\\mathcal{F}(x) + x$.\n",
    "##### Parameters\n",
    "- `fn` - module of type `needle.nn.Module`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_nn_residual\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"nn_residual\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Question 3\n",
    "\n",
    "Implement the `step` function of the following optimizers in `python/needle/optim.py`.\n",
    "Make sure that your optimizers _don't_ modify the gradients of tensors in-place.\n",
    "\n",
    "We have included some tests to ensure that you are not consuming excessive memory, which can happen if you are\n",
    "not using `.data` or `.detach()` in the right places, thus building an increasingly large computational graph\n",
    "(not just in the optimizers, but in the previous modules as well).\n",
    "You can ignore these tests, which include the string `memory_check` at your own discretion.\n",
    "\n",
    "___\n",
    "\n",
    "### SGD\n",
    "`needle.optim.SGD(params, lr=0.01, momentum=0.0, weight_decay=0.0)`\n",
    "\n",
    "Implements stochastic gradient descent (optionally with momentum, shown as $\\beta$ below). \n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{split}\n",
    "    u_{t+1} &= \\beta u_t + (1-\\beta) \\nabla_\\theta f(\\theta_t) \\\\\n",
    "    \\theta_{t+1} &= \\theta_t - \\alpha u_{t+1}\n",
    "\\end{split}\n",
    "\\end{equation}\n",
    "\n",
    "##### Parameters\n",
    "- `params` - iterable of parameters of type `needle.nn.Parameter` to optimize\n",
    "- `lr` (*float*) - learning rate\n",
    "- `momentum` (*float*) - momentum factor\n",
    "- `weight_decay` (*float*) - weight decay (L2 penalty)\n",
    "___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_optim_sgd\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"optim_sgd\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Adam\n",
    "`needle.optim.Adam(params, lr=0.01, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0)`\n",
    "\n",
    "Implements Adam algorithm, proposed in [Adam: A Method for Stochastic Optimization](https://arxiv.org/abs/1412.6980). \n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{split}\n",
    "u_{t+1} &= \\beta_1 u_t + (1-\\beta_1) \\nabla_\\theta f(\\theta_t) \\\\\n",
    "v_{t+1} &= \\beta_2 v_t + (1-\\beta_2) (\\nabla_\\theta f(\\theta_t))^2 \\\\\n",
    "\\hat{u}_{t+1} &= u_{t+1} / (1 - \\beta_1^t) \\quad \\text{(bias correction)} \\\\\n",
    "\\hat{v}_{t+1} &= v_{t+1} / (1 - \\beta_2^t) \\quad \\text{(bias correction)}\\\\\n",
    "\\theta_{t+1} &= \\theta_t - \\alpha \\hat{u_{t+1}}/(\\hat{v}_{t+1}^{1/2}+\\epsilon)\n",
    "\\end{split}\n",
    "    \\end{equation}\n",
    "\n",
    "**Important:** Pay attention to whether or not you are applying bias correction.\n",
    "\n",
    "##### Parameters\n",
    "- `params` - iterable of parameters of type `needle.nn.Parameter` to optimize\n",
    "- `lr` (*float*) - learning rate\n",
    "- `beta1` (*float*) - coefficient used for computing running average of gradient\n",
    "- `beta2` (*float*) - coefficient used for computing running average of square of gradient\n",
    "- `eps` (*float*) - term added to the denominator to improve numerical stability\n",
    "- `weight_decay` (*float*) - weight decay (L2 penalty)\n",
    "\n",
    "**Hint**: To help deal with memory issues, try to understand how to use `.data` or `.detach()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_optim_adam\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"optim_adam\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 4\n",
    "\n",
    "In this question, you will implement two data primitives: `needle.data.DataLoader` and `needle.data.Dataset`. `Dataset` stores the samples and their corresponding labels, and `DataLoader` wraps an iterable around the `Dataset` to enable easy access to the samples. \n",
    "\n",
    "For this question, you will be working in the `python/needle/data` directory. \n",
    "\n",
    "### Transformations\n",
    "\n",
    "First we will implement a few transformations that are helpful when working with images. We will stick with a horizontal flip and a random crop for now. Fill out the following functions in `needle/data/data_transforms.py`.\n",
    "___ \n",
    "\n",
    "#### RandomFlipHorizontal\n",
    "`needle.data.RandomFlipHorizontal(p = 0.5)`\n",
    "\n",
    "Flips the image horizontally, with probability `p`.\n",
    "\n",
    "##### Parameters\n",
    "- `p` (*float*) - The probability of flipping the input image.\n",
    "___\n",
    "\n",
    "#### RandomCrop\n",
    "`needle.data.RandomCrop(padding=3)`\n",
    "\n",
    "Padding is added to all sides of the image, and then the image is cropped back to it's original size at a random location. Returns an image the same size as the original image.\n",
    "\n",
    "##### Parameters\n",
    "- `padding` (*int*) - The padding on each border of the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"flip_horizontal\"\n",
    "!python3 -m pytest -v -k \"random_crop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"flip_horizontal\"\n",
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"random_crop\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "\n",
    "Each `Dataset` subclass must implement three functions: `__init__`, `__len__`, and `__getitem__`. The `__init__` function initializes the images, labels, and transforms. The `__len__` function returns the number of samples in the dataset. The `__getitem__` function retrieves a sample from the dataset at a given index `idx`, calls the transform functions on the image (if applicable), converts the image and label to a numpy array (the data will be converted to Tensors elsewhere). The output of `__getitem__` and `__next__` should be NDArrays, and you should follow the shapes such that you're accessing an array of size (Datapoint Number, Feature Dim 1, Feature Dim 2, ...). \n",
    "\n",
    "Fill out these functions in the `MNISTDataset` class in `needle/data/datasets/mnist_dataset.py`. You can use your solution to `parse_mnist` from the previous homework for the `__init__` function.\n",
    "\n",
    "### MNISTDataset\n",
    "`needle.data.MNISTDataset(image_filesname, label_filesname, transforms)`\n",
    "\n",
    "##### Parameters\n",
    "- `image_filesname` - path of file containing images\n",
    "- `label_filesname` - path of file containing labels\n",
    "- `transforms` - an optional list of transforms to apply to data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_mnist_dataset\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"mnist_dataset\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataloader\n",
    "\n",
    "In `needle/data/data_basic.py`, the Dataloader class provides an interface for assembling mini-batches of examples suitable for training using SGD-based approaches, backed by a Dataset object.  In order to build the typical Dataloader interface (allowing users to iterate over all the mini-batches in the dataset), you will need to implement the `__iter__()` and `__next__()` calls in the class: `__iter__()` is called at the start of iteration, while `__next__()` is called to grab the next mini-batch. Please note that subsequent calls to next will require you to return the following batches, so next is not a pure function.\n",
    "\n",
    "### Dataloader\n",
    "`needle.data.Dataloader(dataset: Dataset, batch_size: Optional[int] = 1, shuffle: bool = False)`\n",
    "\n",
    "Combines a dataset and a sampler, and provides an iterable over the given dataset. \n",
    "\n",
    "##### Parameters\n",
    "- `dataset` - `needle.data.Dataset` - a dataset \n",
    "- `batch_size` - `int` - what batch size to serve the data in \n",
    "- `shuffle` - `bool` - set to ``True`` to have the data reshuffle at every epoch, default ``False``.\n",
    "___ \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_dataloader\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"dataloader\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 5\n",
    "\n",
    "Given you have now implemented all the necessary components for our neural network library, let's build and train an MLP ResNet. For this question, you will be working in `apps/mlp_resnet.py`. First, fill out the functions `ResidualBlock` and `MLPResNet` as described below:\n",
    "\n",
    "### ResidualBlock\n",
    "`ResidualBlock(dim, hidden_dim, norm=nn.BatchNorm1d, drop_prob=0.1)`\n",
    "\n",
    "Implements a residual block as follows:\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"https://raw.github.com/dlsyscourse/hw2/blob/main/figures/residualblock.png\" alt=\"Residual Block\"/>\n",
    "</p>\n",
    "\n",
    "**NOTE**: if the figure does not render, please see the figure in the `figures` directory.\n",
    "\n",
    "where the first linear layer has `in_features=dim` and `out_features=hidden_dim`, and the last linear layer has `out_features=dim`. Returns the block as type `nn.Module`. \n",
    "\n",
    "##### Parameters\n",
    "- `dim` (*int*) - input dim\n",
    "- `hidden_dim` (*int*) - hidden dim\n",
    "- `norm` (*nn.Module*) - normalization method\n",
    "- `drop_prob` (*float*) - dropout probability\n",
    "\n",
    "___\n",
    "\n",
    "### MLPResNet\n",
    "`MLPResNet(dim, hidden_dim=100, num_blocks=3, num_classes=10, norm=nn.BatchNorm1d, drop_prob=0.1)`\n",
    "\n",
    "Implements an MLP ResNet as follows:\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"https://raw.github.com/dlsyscourse/hw2/blob/main/figures/mlp_resnet.png\" alt=\"MLP Resnet\"/>\n",
    "</p>\n",
    "\n",
    "where the first linear layer has `in_features=dim` and `out_features=hidden_dim`, and each ResidualBlock has `dim=hidden_dim` and `hidden_dim=hidden_dim//2`. Returns a network of type `nn.Module`.\n",
    "\n",
    "##### Parameters\n",
    "- `dim` (*int*) - input dim\n",
    "- `hidden_dim` (*int*) - hidden dim\n",
    "- `num_blocks` (*int*) - number of ResidualBlocks\n",
    "- `num_classes` (*int*) - number of classes\n",
    "- `norm` (*nn.Module*) - normalization method\n",
    "- `drop_prob` (*float*) - dropout probability (0.1)\n",
    "\n",
    "**Note**: Modules should be initialized to match the order of execution in the Resnet.\n",
    "___ \n",
    "\n",
    "Once you have the deep learning model architecture correct, let's train the network using our new neural network library components. Specifically, implement the functions `epoch` and `train_mnist`.\n",
    "\n",
    "### Epoch\n",
    "\n",
    "`epoch(dataloader, model, opt=None)`\n",
    "\n",
    "Executes one epoch of training or evaluation, iterating over the entire training dataset once (just like `nn_epoch` from previous homeworks). Returns the average error rate (as a *float*) and the average loss over all samples (as a *float*). Set the model to `training` mode at the beginning of the function if `opt` is given; set the model to `eval` if `opt` is not given (i.e. `None`). When setting the modes, use `.train()` and `.eval()` instead of modifying the training attribute.\n",
    "\n",
    "##### Parameters\n",
    "- `dataloader` (*`needle.data.DataLoader`*) - dataloader returning samples from the training dataset\n",
    "- `model` (*`needle.nn.Module`*) - neural network\n",
    "- `opt` (*`needle.optim.Optimizer`*) - optimizer instance, or `None`\n",
    "\n",
    "___\n",
    "\n",
    "### Train Mnist\n",
    "\n",
    "`train_mnist(batch_size=100, epochs=10, optimizer=ndl.optim.Adam, lr=0.001, weight_decay=0.001, hidden_dim=100, data_dir=\"data\")`\n",
    "                \n",
    "Initializes a training dataloader (with `shuffle` set to `True`) and a test dataloader for MNIST data, and trains an `MLPResNet` using the given optimizer (if `opt` is not None) and the softmax loss for a given number of epochs. Returns a tuple of the training error, training loss, test error, test loss computed in the last epoch of training. If any parameters are not specified, use the default parameters.\n",
    "\n",
    "##### Parameters\n",
    "- `batch_size` (*int*) - batch size to use for train and test dataloader\n",
    "- `epochs` (*int*) - number of epochs to train for\n",
    "- `optimizer` (*`needle.optim.Optimizer` type*) - optimizer type to use\n",
    "- `lr` (*float*) - learning rate \n",
    "- `weight_decay` (*float*) - weight decay\n",
    "- `hidden_dim` (*int*) - hidden dim for `MLPResNet`\n",
    "- `data_dir` (*int*) - directory containing MNIST image/label files\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -v -k \"test_mlp\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m mugrade submit 'YOUR_GRADER_KEY_HERE' -k \"mlp_resnet\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We encourage to experiment with the `mlp_resnet.py` training script.\n",
    "You can investigate the effect of using different initializers on the Linear layers,\n",
    "increasing the dropout probability,\n",
    "or adding transforms (via a list to the `transforms=` keyword argument of Dataset)\n",
    "such as random cropping."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}