Assignment_Implementing_An_MLP_From_Scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Assignment_Implementing_An_MLP_From_Scratch.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4EllkTsH7x0A",
        "colab_type": "text"
      },
      "source": [
        "You will write your own implementation of the backpropagation algorithm for training your own neural network, as\n",
        "well as a few other features such as activation and loss functions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YizmHKDD75-p",
        "colab_type": "text"
      },
      "source": [
        "## Task 1: Activations\n",
        "\n",
        "Implement the `forward` and `derivative` class methods for each activation function.\n",
        "* The identity function has been implemented for you as an example.\n",
        "* The output of the activation should be stored in the `self.state` variable of the class. The `self.state`\n",
        "variable should be used for calculating the derivative during the backward pass."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sHf0B0FQ8aOT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import math"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZUlNOZKH6xD0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Activation(object):\n",
        "\n",
        "    \"\"\"\n",
        "    Interface for activation functions (non-linearities).\n",
        "\n",
        "    In all implementations, the state attribute must contain the result,\n",
        "    i.e. the output of forward.\n",
        "    \"\"\"\n",
        "\n",
        "    # No additional work is needed for this class, as it acts like an\n",
        "    # abstract base class for the others\n",
        "\n",
        "    # Note that these activation functions are scalar operations. I.e, they\n",
        "    # shouldn't change the shape of the input.\n",
        "\n",
        "    def __init__(self):\n",
        "        self.state = None\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return self.forward(x)\n",
        "\n",
        "    def forward(self, x):\n",
        "        raise NotImplemented\n",
        "\n",
        "    def derivative(self):\n",
        "        raise NotImplemented"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s42anq5n8X1b",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Identity(Activation):\n",
        "\n",
        "    \"\"\"\n",
        "    Identity function (already implemented).\n",
        "    \"\"\"\n",
        "\n",
        "    # This class is a gimme as it is already implemented for you as an example\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Identity, self).__init__()\n",
        "\n",
        "    def forward(self, x):\n",
        "        self.state = x\n",
        "        return x\n",
        "    # Design Choices\n",
        "    def derivative(self):\n",
        "        return 1.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4ykhtBPg8iDl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Sigmoid(Activation):\n",
        "\n",
        "    \"\"\"\n",
        "    Sigmoid non-linearity\n",
        "    \"\"\"\n",
        "\n",
        "    # Remember do not change the function signatures as those are needed\n",
        "    # to stay the same for AutoLab.\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Sigmoid, self).__init__()\n",
        "    # vectorization\n",
        "    def forward(self, x):\n",
        "        # Might we need to store something before returning?\n",
        "        self.state=1.0/(1.0 + exp(-x))\n",
        "        return self.state\n",
        "\n",
        "    def derivative(self):\n",
        "        # Maybe something we need later in here...\n",
        "        return self.state * (1-self.state)\n",
        "        "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-t_MXcXU8l0i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Tanh(Activation):\n",
        "\n",
        "    \"\"\"\n",
        "    Tanh non-linearity\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Tanh, self).__init__()\n",
        "\n",
        "    def forward(self, x):\n",
        "        e_minus_x=np.exe(-x)\n",
        "        e_plus_x=np.exe(x)\n",
        "\n",
        "        numerator = e_plus_x + e_minus_x\n",
        "        denominator = e_plus_x - e_minus_x\n",
        "        self.state = numerator/denominator\n",
        "        return self.state\n",
        "\n",
        "    def derivative(self):\n",
        "        return 1 - (self.state * self.state)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KGVifuPH8n5Z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ReLU(Activation):\n",
        "\n",
        "    \"\"\"\n",
        "    ReLU non-linearity\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(ReLU, self).__init__()\n",
        "\n",
        "    def forward(self, x):\n",
        "        self.state = np.where(x>0, x, 0.0)\n",
        "        return self.state\n",
        "\n",
        "    def derivative(self):\n",
        "        relu_prime = np.where(self.state<0, 1.0, 0.0)\n",
        "        return relu_prime"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LD-Qch_F8vPU",
        "colab_type": "text"
      },
      "source": [
        "## Task 2: Loss\n",
        "Implement the forward and derivative methods for `SoftmaxCrossEntropy`.\n",
        "* This class inherits the base `Criterion` class.\n",
        "* We will be using the softmax cross entropy loss detailed in the appendix of this writeup; use the\n",
        "LogSumExp trick to ensure numerical stability.\n",
        "\n",
        "The LogSumExp trick is used to prevent numerical underflow and overflow which can occur when the exponent is very large or very small. For example, try looking at the results of trying to exponentiate in python shown below:\n",
        "\n",
        "```python\n",
        "import math\n",
        "print(math.e**1000)  # throws an error\n",
        "print(math.e**(-1000)\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "As you will see, for exponents that are too large, python throws an overflow error, and for exponents that are too small, it rounds down to zero.\n",
        "We can avoid these errors by using the LogSumExp trick:\n",
        "\n",
        "![alt text](https://imgur.com/download/L0P17iv)\n",
        "\n",
        "You can read more about the derivation of the equivalence [here](https://www.xarg.org/2016/06/the-log-sum-exp-trick-in-machine-learning/) and [here](https://blog.feedly.com/tricks-of-the-trade-logsumexp/)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r4CsuOXy88eo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# The following Criterion class will be used again as the basis for a number\n",
        "# of loss functions (which are in the form of classes so that they can be\n",
        "# exchanged easily (it's how PyTorch and other ML libraries do it))\n",
        "\n",
        "class Criterion(object):\n",
        "    \"\"\"\n",
        "    Interface for loss functions.\n",
        "    \"\"\"\n",
        "\n",
        "    # Nothing needs done to this class, it's used by the following Criterion classes\n",
        "\n",
        "    def __init__(self):\n",
        "        self.logits = None\n",
        "        self.labels = None\n",
        "        self.loss = None\n",
        "\n",
        "    def __call__(self, x, y):\n",
        "        return self.forward(x, y)\n",
        "\n",
        "    def forward(self, x, y):\n",
        "        raise NotImplemented\n",
        "\n",
        "    def derivative(self):\n",
        "        raise NotImplemented"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFJezVji9Opc",
        "colab_type": "text"
      },
      "source": [
        "* Implement the softmax cross entropy operation on a batch of output vectors.\n",
        "  *  Hint: Add a class attribute to keep track of intermediate values necessary for the backward computation\n",
        "* Calculate the ‘derivative’ of softmax cross entropy using intermediate values saved in the forward pass."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nijZM_m09HU6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SoftmaxCrossEntropy(Criterion):\n",
        "    \"\"\"\n",
        "    Softmax loss\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(SoftmaxCrossEntropy, self).__init__()\n",
        "\n",
        "    def forward(self, x, y):\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            x (np.array): (batch size, 10)\n",
        "            y (np.array): (batch size, 10)\n",
        "        Return:\n",
        "            out (np.array): (batch size, )\n",
        "        \"\"\"\n",
        "        self.logits = x\n",
        "        self.labels = y\n",
        "\n",
        "        # Numerically stable softmax\n",
        "        mx=np.max(self.logits, axis =1).reshape(-1,1)\n",
        "        subtracted = self.logits - mx\n",
        "        self.exp_logits = np.exp(subtracted)\n",
        "        exp_sum = self.exp_logits.sum(axis=1).reshape(-1,1)\n",
        "        self.sm = self.exp_logits / exp_sum\n",
        "\n",
        "        # cross entropy\n",
        "        first_term = -(self.logits*self.labels).sum(axis = 1)\n",
        "        second_term = mx + np.log(exp_sum)\n",
        "        return first_term + second_term.reshape(-1)\n",
        "\n",
        "    def derivative(self):\n",
        "        \"\"\"\n",
        "        Return:\n",
        "            out (np.array): (batch size, 10)\n",
        "        \"\"\"\n",
        "        return self.sm - self.labels"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MdolnMqE9idk",
        "colab_type": "text"
      },
      "source": [
        "## Task 3: Linear Layer\n",
        "Implement the forward and backward methods for the `Linear` class.\n",
        "* Hint: Add a class attribute to keep track of intermediate values necessary for the backward computation.\n",
        "\n",
        "Write the code for the backward method of Linear. \n",
        "* The input delta is the derivative of the loss with respect to the output of the linear layer. It has the same shape as the linear layer output. \n",
        "* Calculate `self.dW` and `self.db` for the backward method. `self.dW` and `self.db` represent the gradients of the loss (averaged across the batch) w.r.t `self.W` and `self.b`. Their shapes are the same as the weight `self.W` and the bias `self.b`.\n",
        "* Calculate the return value for the backward method. `dx` is the derivative of the loss with respect to the input of the linear layer and has the same shape as the input."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TU_EST6O9J34",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Weights as metrices \n",
        "\n",
        "class Linear():\n",
        "    def __init__(self, in_feature, out_feature, weight_init_fn, bias_init_fn):\n",
        "\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            W (np.array): (in feature, out feature)\n",
        "            dW (np.array): (in feature, out feature)\n",
        "            momentum_W (np.array): (in feature, out feature)\n",
        "\n",
        "            b (np.array): (1, out feature)\n",
        "            db (np.array): (1, out feature)\n",
        "            momentum_B (np.array): (1, out feature)\n",
        "        \"\"\"\n",
        "        # Design Choices:\n",
        "        self.W = weight_init_fn(in_feature, out_feature)\n",
        "        self.b = bias_init_fn(out_feature)\n",
        "\n",
        "        # Complete these but do not change the names.\n",
        "        self.dW = np.zeros(None)\n",
        "        self.db = np.zeros(None)\n",
        "\n",
        "        self.momentum_W = np.zeros(None)\n",
        "        self.momentum_b = np.zeros(None)\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return self.forward(x)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            delta (=delL/delY) (np.array): (batch size, in feature)\n",
        "            W: 784, 10\n",
        "            x: 20, 784\n",
        "            out = 20, 10\n",
        "            delta: 20, 10\n",
        "        Return:\n",
        "            out (np.array): (batch size, out feature)\n",
        "        \"\"\"\n",
        "        self.x = x\n",
        "        # Batch Processing\n",
        "        out = np.matmul(self.x, self.W) + self.backward\n",
        "        return out\n",
        "\n",
        "    def backward(self, delta):\n",
        "\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            delta (np.array): (batch size, out feature)\n",
        "        Return:\n",
        "            out (np.array): (batch size, in feature)\n",
        "        \"\"\"\n",
        "        # Point out strong gradients to avoid recomputation (W.T= transpose(W))\n",
        "        self.dW = np.dot(self.x.T, delta)/dela.shape(0)\n",
        "        self.db = np.sum(delta, axis=0, keepdims=True) / delta.shape(0)\n",
        "        dx = np.dot(delta, self.W.T)\n",
        "        return dx\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gWkv9PDr-Wcy",
        "colab_type": "text"
      },
      "source": [
        "## Task 4: Simple MLP\n",
        "In this section of the homework, you will be implementing a Multi-Layer  Perceptron with an API similar to popular Automatic Differentiation Libraries like PyTorch.\n",
        "Go through the functions of the given `MLP` class thoroughly and make sure you understand what each function in the class does so that you can create a generic implementation that supports an arbitrary number of layers, types of activations and network sizes.\n",
        "\n",
        "The parameters for the MLP class are:\n",
        "* `input size`: The size of each individual data example.\n",
        "* `output size`: The number of outputs.\n",
        "* `hiddens`: A list with the number of units in each hidden layer.\n",
        "* `activations`: A list of Activation objects for each layer.\n",
        "* `weight init fn`: A function applied to each weight matrix before training.\n",
        "* `bias init fn`: A function applied to each bias vector before training.\n",
        "* `criterion`: A Criterion object to compute the loss and its derivative.\n",
        "* `lr`: The learning rate.\n",
        "\n",
        "The attributes of the MLP class are:\n",
        "* `@linear layers`: A list of Linear objects.\n",
        "* `@bn layers`: A list of BatchNorm objects. (Should be None until completing 3.3).\n",
        "The methods of the MLP class are:\n",
        "* `forward`: Forward pass. Accepts a mini-batch of data and return a batch of output activations.\n",
        "* `backward`: Backward pass. Accepts ground truth labels and computes gradients for all parameters.\n",
        "Hint: Use state stored in activations during forward pass to simplify your code.\n",
        "* `zero grads`: Set all gradient terms to 0.\n",
        "* `step`: Apply gradients computed in backward to the parameters.\n",
        "* `train` (Already implemented): Set the mode of the network to train.\n",
        "* `eval` (Already implemented): Set the mode of the network to evaluation.\n",
        "\n",
        "Note: Pay attention to the data structures being passed into the constructor and the class attributes specified initially.\n",
        "\n",
        "Sample constructor call:\n",
        "```python\n",
        "MLP(784, 10, [64, 64, 32], [Sigmoid(), Sigmoid(), Sigmoid(), Identity()],\n",
        "weight_init_fn, bias_init_fn, SoftmaxCrossEntropy(), 0.008)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0pIxdtci-R4S",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MLP(object):\n",
        "\n",
        "    \"\"\"\n",
        "    A simple multilayer perceptron\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, input_size, output_size, hiddens, activations, weight_init_fn,\n",
        "                 bias_init_fn, criterion, lr):\n",
        "\n",
        "        # Don't change this -->\n",
        "        self.train_mode = True\n",
        "        self.nlayers = len(hiddens) + 1\n",
        "        self.input_size = input_size\n",
        "        self.output_size = output_size\n",
        "        self.activations = activations\n",
        "        self.criterion = criterion\n",
        "        self.lr = lr\n",
        "        # <---------------------\n",
        "\n",
        "        # Don't change the name of the following class attributes,\n",
        "        # the autograder will check against these attributes. But you will need to change\n",
        "        # the values in order to initialize them correctly\n",
        "\n",
        "        # Initialize and add all your linear layers into the list 'self.linear_layers'\n",
        "        # (HINT: self.foo = [ bar(???) for ?? in ? ])\n",
        "        # (HINT: Can you use zip here?)\n",
        "        self.linear_layers = (Linear(inf, outf, weight_init_fn, bias_init_fn) \n",
        "                              for inf, outf in zip ([self.input_size] + hidens, hiddens + [self.output_size]))\n",
        "\n",
        "          \n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            x (np.array): (batch size, input_size)\n",
        "        Return:\n",
        "            out (np.array): (batch size, output_size)\n",
        "        \"\"\"\n",
        "        # Complete the forward pass through your entire MLP.\n",
        "        for i, layer in enumerate(self.linear_layers):\n",
        "          x = layer(x)\n",
        "          x = self.activations[i](x)\n",
        "        return x\n",
        "\n",
        "    def zero_grads(self):\n",
        "        # Use numpyArray.fill(0.0) to zero out your backpropped derivatives in each\n",
        "        # of your linear and batchnorm layers.\n",
        "        for layer in self.linear_layers:\n",
        "          layer.dW.fill(0.0)\n",
        "          layer.db.fill(0.0)\n",
        "       \n",
        "    def step(self):\n",
        "        # Step function:: Apply a step to the weights and biases of the linear layers.\n",
        "        # (You will add momentum later in the assignment to the linear layers)\n",
        "\n",
        "        # without momentum\n",
        "        for i in range(len(self.linear_layers)):\n",
        "            # Update weights and biases here\n",
        "            layer = self.linear_layers[i]\n",
        "            layer.W = layer.W - self.lr * layer.dW\n",
        "            layer.b = layer.b - self.lr * layer.db\n",
        "            \n",
        "    def backward(self, labels):\n",
        "        # Backpropagate through the activation functions, batch norm and\n",
        "        # linear layers.\n",
        "        # Be aware of which return derivatives and which are pure backward passes\n",
        "        # i.e. take in a loss w.r.t it's output.\n",
        "        final_layer = self.activations[-1]\n",
        "        final_outputs = final_layer.state\n",
        "        loss = self.criterion(final_outputs, labels)\n",
        "        delta = self.criterion.derivative()\n",
        "\n",
        "        # Using sigmoid complicate the computation. Design Choices::\n",
        "        for i in range(self.nlayers - 1, -1, -1):\n",
        "          delta = delta + self.activations[i].derivative()\n",
        "          delta = self.linear_layers(i).backward(delta)\n",
        "\n",
        "    def error(self, labels):\n",
        "        return (np.argmax(self.output, axis = 1) != np.argmax(labels, axis = 1)).sum()\n",
        "\n",
        "    def total_loss(self, labels):\n",
        "        return self.criterion(self.output, labels).sum()\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return self.forward(x)\n",
        "\n",
        "    def train(self):\n",
        "        self.train_mode = True\n",
        "\n",
        "    def eval(self):\n",
        "        self.train_mode = False\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yMrjk_-FAQfZ",
        "colab_type": "text"
      },
      "source": [
        "## Task 5: Momentum\n",
        "Modify the `step` function present in the MLP class to include momentum in your gradient descent. We will be using the following momentum update equation:\n",
        "\n",
        "![alt text](https://imgur.com/download/ZVA66FC)\n",
        "\n",
        "The momentum value will be passed as a parameter to the `MLP`.\n",
        "Copy the rest of your code from above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SUZJiCnoANLL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MLP(object):\n",
        "\n",
        "    \"\"\"\n",
        "    A simple multilayer perceptron\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, input_size, output_size, hiddens, activations, weight_init_fn,\n",
        "                 bias_init_fn, criterion, lr, momentum=0.0):\n",
        "\n",
        "        # Don't change this -->\n",
        "        self.train_mode = True\n",
        "        self.nlayers = len(hiddens) + 1\n",
        "        self.input_size = input_size\n",
        "        self.output_size = output_size\n",
        "        self.activations = activations\n",
        "        self.criterion = criterion\n",
        "        self.lr = lr\n",
        "        self.momentum = momentum\n",
        "        # <---------------------\n",
        "\n",
        "        # Initialize and add all your linear layers into the list 'self.linear_layers'\n",
        "        # (HINT: self.foo = [ bar(???) for ?? in ? ])\n",
        "        # (HINT: Can you use zip here?)\n",
        "        self.linear_layers = (Linear(inf, outf, weight_init_fn, bias_init_fn) \n",
        "                              for inf, outf in zip ([self.input_size] + hidens, hiddens + [self.output_size]))\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Argument:\n",
        "            x (np.array): (batch size, input_size)\n",
        "        Return:\n",
        "            out (np.array): (batch size, output_size)\n",
        "        \"\"\"\n",
        "        # Complete the forward pass through your entire MLP.\n",
        "        for i, layer in enumerate(self.linear):\n",
        "          x = layer(x)\n",
        "          x = self.activations[i](x)\n",
        "        return x\n",
        "\n",
        "    def zero_grads(self):\n",
        "        # Use numpyArray.fill(0.0) to zero out your backpropped derivatives in each\n",
        "        # of your linear and batchnorm layers.\n",
        "        for layer in self.linear_layers:\n",
        "          layer.dW.fill(0.0)\n",
        "          layer.db.fill(0.0)\n",
        "          \n",
        "    def step(self):\n",
        "        # Apply a step to the weights and biases of the linear layers.\n",
        "        # (You will add momentum later in the assignment to the linear layers)\n",
        "\n",
        "        for i in range(len(self.linear_layers)):\n",
        "            # Update weights and biases here\n",
        "            layer = self.linear_layers[i]\n",
        "            layer.momentum_W = self.momentum * layer.momentum_W - self.lr * layer.dW\n",
        "            layer.W = layer.W + layer.momentum_W\n",
        "            layer.momentum_b = self.momentum * layer.momentum_b - self.lr * layer.db\n",
        "            layer.b = layer.b + layer.momentum_b\n",
        "        # Do the same for batchnorm layers\n",
        "\n",
        "    def backward(self, labels):\n",
        "        # Backpropagate through the activation functions, batch norm and\n",
        "        # linear layers.\n",
        "        # Be aware of which return derivatives and which are pure backward passes\n",
        "        # i.e. take in a loss w.r.t it's output.\n",
        "        final_layer = self.activations[-1]\n",
        "        final_outputs = final_layer.state\n",
        "        loss = self.criterion(final_outputs, labels)\n",
        "        delta = self.criterion.derivative()\n",
        "\n",
        "        for i in range(self.nlayers-1, -1, -1):\n",
        "          delta = delta * self.activation[i].derivative()\n",
        "          delta = self.linear_layers[i].backward(delta)\n",
        "\n",
        "    def error(self, labels):\n",
        "        return (np.argmax(self.output, axis = 1) != np.argmax(labels, axis = 1)).sum()\n",
        "\n",
        "    def total_loss(self, labels):\n",
        "        return self.criterion(self.output, labels).sum()\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return self.forward(x)\n",
        "\n",
        "    def train(self):\n",
        "        self.train_mode = True\n",
        "\n",
        "    def eval(self):\n",
        "        self.train_mode = False\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}