{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyze (F)MNIST with `torch`\n",
    "\n",
    "Careful: do **not** hit 'Kernel' > 'Restart & Run All', since some of the cells below take a long time to execute if you are not running the code on a GPU, so we already executed them for you. Only run the first few cells that are not yet executed.\n",
    "\n",
    "In this notebook we compare different types of neural network architectures on the MNIST and Fashion MNIST datasets, to see how the performance improves when using a more complicated architecture. Additionally, we compare the networks to a simple logistic regression classifier from `sklearn`, which should have approximately the same accuracy as a linear FFNN (= a FFNN with only one layer mapping from the input directly to the output and no hidden layers, i.e., that has the same number of trainable parameters as the logistic regression model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:06:30.088245Z",
     "start_time": "2020-11-22T19:06:29.139733Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import accuracy_score\n",
    "# torch neural network stuff\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "# torchvision includes the (F)MNIST datasets\n",
    "from torchvision import datasets, transforms\n",
    "# skorch provides a wrapper for torch networks so we can use them like sklearn models\n",
    "from skorch import NeuralNetClassifier\n",
    "from skorch.callbacks import EpochScoring\n",
    "# set random seeds to get (at least more or less) reproducable results\n",
    "np.random.seed(28)\n",
    "torch.manual_seed(28);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and look at the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:06:30.114557Z",
     "start_time": "2020-11-22T19:06:30.090761Z"
    }
   },
   "outputs": [],
   "source": [
    "# you do not need to understand what these functions do in detail\n",
    "\n",
    "def torch_to_X_y(dataset):\n",
    "    # transform input tensor to numpy array\n",
    "    X = dataset.data.numpy()\n",
    "    # reshape (28 x 28) pixel images to vector\n",
    "    X = X.reshape(X.shape[0], -1).astype('float32')\n",
    "    # the ToTensor transform was not applied to the raw data, so we need to scale ourselves\n",
    "    X /= X.max()\n",
    "    # extract numpy array with targets\n",
    "    y = dataset.targets.numpy()\n",
    "    return X, y\n",
    "\n",
    "def load_data(use_fashion=False):\n",
    "    if use_fashion:\n",
    "        data_train = datasets.FashionMNIST(\"../data\", train=True, download=True, transform=transforms.ToTensor())\n",
    "        data_test = datasets.FashionMNIST(\"../data\", train=False, transform=transforms.ToTensor())\n",
    "    else:\n",
    "        data_train = datasets.MNIST(\"../data\", train=True, download=True, transform=transforms.ToTensor())\n",
    "        data_test = datasets.MNIST(\"../data\", train=False, transform=transforms.ToTensor())\n",
    "    # extract (n_samples x n_features) and (n_samples,) X and y numpy arrays from torch dataset\n",
    "    X_train, y_train = torch_to_X_y(data_train)\n",
    "    X_test, y_test = torch_to_X_y(data_test)\n",
    "    return X_train, X_test, y_train, y_test\n",
    "    \n",
    "def plot_images(x):\n",
    "    n = 10\n",
    "    plt.figure(figsize=(20, 4))\n",
    "    for i in range(1, n+1):\n",
    "        # display original\n",
    "        ax = plt.subplot(2, n, i)\n",
    "        plt.imshow(x[i].reshape(28, 28))\n",
    "        plt.gray()\n",
    "        ax.get_xaxis().set_visible(False)\n",
    "        ax.get_yaxis().set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:06:31.076712Z",
     "start_time": "2020-11-22T19:06:30.115973Z"
    }
   },
   "outputs": [],
   "source": [
    "# load and display the data -> see how the images have the same format\n",
    "# MNIST\n",
    "X_train, X_test, y_train, y_test = load_data()\n",
    "plot_images(X_train)\n",
    "# Fashion MNIST\n",
    "X_train_F, X_test_F, y_train_F, y_test_F = load_data(use_fashion=True)\n",
    "plot_images(X_train_F)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See how a `torch` network works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:06:31.158130Z",
     "start_time": "2020-11-22T19:06:31.078469Z"
    }
   },
   "outputs": [],
   "source": [
    "# FFNN with hidden layers (like the one you saw in the book)\n",
    "class MyNeuralNet(nn.Module):\n",
    "    \n",
    "    def __init__(self, n_in=784, n_hl1=512, n_hl2=256, n_out=10, verbose=False):\n",
    "        # input size is 28x28 pixel images flattened into a 784-dimensional vector\n",
    "        # output size is 10 classes\n",
    "        # hidden layer sizes can be set however you want\n",
    "        super(MyNeuralNet, self).__init__()\n",
    "        self.verbose = verbose\n",
    "        # initialize layers\n",
    "        self.l1 = nn.Linear(n_in, n_hl1)\n",
    "        self.l2 = nn.Linear(n_hl1, n_hl2)\n",
    "        self.lout = nn.Linear(n_hl2, n_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # apply layers in correct order\n",
    "        if self.verbose: print(\"[MyNeuralNet]  input:\", x.shape)\n",
    "        h = F.relu(self.l1(x))              # 784 -> 512 [relu]\n",
    "        if self.verbose: print(\"[MyNeuralNet] 1st hl:\", h.shape)\n",
    "        h = F.relu(self.l2(h))              # 512 -> 256 [relu]\n",
    "        if self.verbose: print(\"[MyNeuralNet] 2nd hl:\", h.shape)\n",
    "        y = F.softmax(self.lout(h), dim=1)  # 256 -> 10  [softmax]\n",
    "        if self.verbose: print(\"[MyNeuralNet] output:\", y.shape)\n",
    "        return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the network\n",
    "ffnn = MyNeuralNet(verbose=True)\n",
    "# get an input data batch and convert the numpy array \n",
    "# to a torch tensor to use it with the network directly\n",
    "# (skorch later works with the numpy arrays)\n",
    "x = torch.Tensor(X_train[:16])\n",
    "print(x.shape)  # batch size x features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply network to input, i.e., call forward() to generate the prediction\n",
    "y = ffnn(x)  # same as: y = ffnn.forward(x)\n",
    "print(y.shape)  # batch size x classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the image of the first training sample\n",
    "plt.imshow(X_train[0].reshape(28, 28));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# look at the network's output for this first data point\n",
    "# -> since the network wasn't trained yet, the predicted probabilities for all 10 classes are ~0.1\n",
    "# (notice the grad parameter, which indicates that the network kept track of the gradients,\n",
    "# which are needed for later tuning the weights during training)\n",
    "y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wrap torch NN in skorch Classifier and initialize\n",
    "net = NeuralNetClassifier(\n",
    "    MyNeuralNet,  # usually the class itself, not an instantiated object\n",
    "    batch_size=32,  # how many samples are used in each training iteration\n",
    "    optimizer=torch.optim.Adadelta,  # the optimizer (i.e. \"what type\" of gradient descent)\n",
    "    lr=1.,  # learning rate of the optimizer\n",
    "    device=\"cuda\" if torch.cuda.is_available() else \"cpu\",  # train the network on a GPU if available\n",
    "    max_epochs=1,  # for how many epochs to train the network\n",
    "    callbacks=[  # additional stuff that should happen after each epoch, e.g., learning rate scheduler\n",
    "        ('tr_acc', EpochScoring(  # or in this case print the accuracy after every epoch\n",
    "            'accuracy',\n",
    "            lower_is_better=False,\n",
    "            on_train=True,\n",
    "            name='train_acc',\n",
    "        )),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# use simple sklearn-like interface to train the network (for 1 epoch)\n",
    "net.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate predictions for the same samples as above\n",
    "# -> this gives class labels directly like sklearn\n",
    "y = net.predict(X_train[:16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if the prediction (after training) is correct\n",
    "print(\"true class:\", y_train[0])\n",
    "y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can also get the original probabilities (notice the higher value at the index of the true class)\n",
    "y = net.predict_proba(X_train[:16])\n",
    "y[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define NNs for the classification task\n",
    "\n",
    "In the code below, we define 3 different neural network architectures: a linear FFNN, a FFNN with multiple hidden layers, and a CNN, which is an architecture particularly well suited for image classification tasks.\n",
    "\n",
    "You will see that the more complex architectures use an additional operation between layers called `Dropout`. This is a regularization technique used for training neural networks, where a certain percentage of the values in the hidden layer representation of a data point are randomly set to zero. You can think of this as the network suffering from a temporary stroke, which forces the neurons learn redundant representations (i.e., such that one neuron can take over for another neuron that was knocked out), which improves generalization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:06:31.188855Z",
     "start_time": "2020-11-22T19:06:31.159480Z"
    }
   },
   "outputs": [],
   "source": [
    "# linear FFNN (--> same number of parameters as LogReg model)\n",
    "class LinNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, n_in=784, n_out=10):\n",
    "        super(LinNN, self).__init__()\n",
    "        self.l = nn.Linear(n_in, n_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        y = F.softmax(self.l(x), dim=1)  # 784 -> 10 [softmax]\n",
    "        return y\n",
    "    \n",
    "# FFNN with hidden layers  \n",
    "class FFNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, n_in=784, n_hl1=512, n_hl2=256, n_out=10, dropout=0.2):\n",
    "        super(FFNN, self).__init__()\n",
    "        # initialize layers\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.l1 = nn.Linear(n_in, n_hl1)\n",
    "        self.l2 = nn.Linear(n_hl1, n_hl2)\n",
    "        self.lout = nn.Linear(n_hl2, n_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # apply layers in correct order\n",
    "        h = F.relu(self.l1(x))              # 784 -> 512 [relu]\n",
    "        h = self.dropout(h)\n",
    "        h = F.relu(self.l2(h))              # 512 -> 256 [relu]\n",
    "        h = self.dropout(h)\n",
    "        y = F.softmax(self.lout(h), dim=1)  # 256 -> 10  [softmax]\n",
    "        return y\n",
    "    \n",
    "# Convolutional Neural Net    \n",
    "# based on https://github.com/pytorch/examples/blob/master/mnist/main.py\n",
    "class CNN(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # convolutional and pooling layers\n",
    "        h = self.conv1(x)\n",
    "        h = F.relu(h)\n",
    "        h = self.conv2(h)\n",
    "        h = F.relu(h)\n",
    "        h = F.max_pool2d(h, 2)\n",
    "        h = self.dropout1(h)\n",
    "        # flatten the representation and apply FFNN part for the classification\n",
    "        h = torch.flatten(h, 1)\n",
    "        h = self.fc1(h)\n",
    "        h = F.relu(h)\n",
    "        h = self.dropout2(h)\n",
    "        h = self.fc2(h)\n",
    "        y = F.softmax(h, dim=1)\n",
    "        return y\n",
    "\n",
    "# skorch wrapper with fit/predict methods\n",
    "def eval_net(net_module, X_train, y_train, X_test, y_test, max_epochs=1):\n",
    "    print(\"###\", net_module.__name__)\n",
    "    net = NeuralNetClassifier(\n",
    "        net_module,\n",
    "        batch_size=32,\n",
    "        optimizer=torch.optim.Adadelta,\n",
    "        lr=1.,\n",
    "        device=\"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
    "        max_epochs=max_epochs,\n",
    "        callbacks=[\n",
    "            ('tr_acc', EpochScoring(\n",
    "                'accuracy',\n",
    "                lower_is_better=False,\n",
    "                on_train=True,\n",
    "                name='train_acc',\n",
    "            )),\n",
    "        ],\n",
    "    )\n",
    "    net.fit(X_train, y_train)\n",
    "    # evaluate on test set\n",
    "    y_pred = net.predict(X_test)\n",
    "    print('Test accuracy:', accuracy_score(y_test, y_pred), \"\\n\")\n",
    "    return net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test on MNIST dataset\n",
    "\n",
    "As you see below, the simple logistic regression classifier is already very good on this easy task, with a test accuracy of over 92.6%.\n",
    "\n",
    "The linear FFNN has almost the same accuracy (91.0%) as the LogReg model (please note: the NNs were only trained for a single epoch!) and the multi-layer FFNN is already better than the LogReg model (96.3%), while the CNN beats them all (98.2%), which is expected since this architecture is designed for the image classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:07:16.662634Z",
     "start_time": "2020-11-22T19:06:31.190048Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### LogReg\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/franzi/miniconda3/lib/python3.9/site-packages/sklearn/linear_model/_logistic.py:763: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 0.9263 \n",
      "\n",
      "### LinNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.8914\u001b[0m        \u001b[32m0.4044\u001b[0m       \u001b[35m0.9048\u001b[0m        \u001b[31m0.3267\u001b[0m  1.8230\n",
      "Test accuracy: 0.9097 \n",
      "\n",
      "### FFNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.9203\u001b[0m        \u001b[32m0.2596\u001b[0m       \u001b[35m0.9587\u001b[0m        \u001b[31m0.1487\u001b[0m  3.6306\n",
      "Test accuracy: 0.9625 \n",
      "\n",
      "### CNN\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/franzi/miniconda3/lib/python3.9/site-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /pytorch/c10/core/TensorImpl.h:1156.)\n",
      "  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.9251\u001b[0m        \u001b[32m0.2474\u001b[0m       \u001b[35m0.9772\u001b[0m        \u001b[31m0.0786\u001b[0m  7.1728\n",
      "Test accuracy: 0.9824 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# get regular MNIST dataset\n",
    "X_train, X_test, y_train, y_test = load_data()\n",
    "# compare sklearn LogReg classifier\n",
    "print(\"### LogReg\")\n",
    "clf = LogisticRegression(class_weight='balanced', random_state=1, fit_intercept=True)\n",
    "clf.fit(X_train, y_train)\n",
    "print('Test accuracy:', clf.score(X_test, y_test), \"\\n\")\n",
    "# and our different NN architectures\n",
    "for net_module in [LinNN, FFNN, CNN]:\n",
    "    if net_module == CNN:\n",
    "        # the CNN operates on the 28x28 pixel images directly\n",
    "        net = eval_net(net_module, X_train.reshape(-1, 1, 28, 28), y_train, X_test.reshape(-1, 1, 28, 28), y_test)\n",
    "    else:\n",
    "        # the FFNNs get the flattened vectors\n",
    "        net = eval_net(net_module, X_train, y_train, X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test on FashionMNIST\n",
    "\n",
    "On the more difficult FMNIST task, the LogReg model has a much lower accuracy of 84.4% compared to the 92.6% achieved on the original MNIST dataset. When trained for only a single epoch, both the linear and multi-layer FFNNs have a lower accuracy than the LogReg model (81.2 and 82.4% respectively) and only the CNN does a bit better (87.3%). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:07:58.876986Z",
     "start_time": "2020-11-22T19:07:16.665145Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### LogReg\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/franzi/miniconda3/lib/python3.9/site-packages/sklearn/linear_model/_logistic.py:763: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 0.8438 \n",
      "\n",
      "### LinNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.8015\u001b[0m        \u001b[32m0.5792\u001b[0m       \u001b[35m0.8253\u001b[0m        \u001b[31m0.5040\u001b[0m  1.9656\n",
      "Test accuracy: 0.8117 \n",
      "\n",
      "### FFNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.7827\u001b[0m        \u001b[32m0.5929\u001b[0m       \u001b[35m0.8400\u001b[0m        \u001b[31m0.4453\u001b[0m  3.8960\n",
      "Test accuracy: 0.8236 \n",
      "\n",
      "### CNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.8000\u001b[0m        \u001b[32m0.5616\u001b[0m       \u001b[35m0.8818\u001b[0m        \u001b[31m0.3272\u001b[0m  7.0208\n",
      "Test accuracy: 0.8728 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = load_data(True)\n",
    "# regular sklearn LogReg classifier\n",
    "print(\"### LogReg\")\n",
    "clf = LogisticRegression(class_weight='balanced', random_state=1, fit_intercept=True)\n",
    "clf.fit(X_train, y_train)\n",
    "print('Test accuracy:', clf.score(X_test, y_test), \"\\n\")\n",
    "# our different NN\n",
    "for net_module in [LinNN, FFNN, CNN]:\n",
    "    if net_module == CNN:\n",
    "        net = eval_net(net_module, X_train.reshape(-1, 1, 28, 28), y_train, X_test.reshape(-1, 1, 28, 28), y_test)\n",
    "    else:\n",
    "        net = eval_net(net_module, X_train, y_train, X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, when trained for more epochs, the performance of all models improves, with the accuracy of the linear FFNN now being very close to that of the LogReg model (84.4%), while the multi-layer FFNN is better (87.1%) and the CNN can now solve the task quite well with an accuracy of 91.7%.\n",
    "\n",
    "(See how the training and validation loss decrease over time - observing how these metrics develop can help you judge whether you've set your learning rate correctly and for how many epochs you should train the network.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T19:12:29.545463Z",
     "start_time": "2020-11-22T19:07:58.880264Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### LinNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.8002\u001b[0m        \u001b[32m0.5800\u001b[0m       \u001b[35m0.8263\u001b[0m        \u001b[31m0.5020\u001b[0m  2.0403\n",
      "      2       \u001b[36m0.8389\u001b[0m        \u001b[32m0.4715\u001b[0m       \u001b[35m0.8372\u001b[0m        \u001b[31m0.4734\u001b[0m  2.1028\n",
      "      3       \u001b[36m0.8458\u001b[0m        \u001b[32m0.4520\u001b[0m       \u001b[35m0.8409\u001b[0m        \u001b[31m0.4620\u001b[0m  2.1199\n",
      "      4       \u001b[36m0.8498\u001b[0m        \u001b[32m0.4416\u001b[0m       \u001b[35m0.8427\u001b[0m        \u001b[31m0.4563\u001b[0m  2.2367\n",
      "      5       \u001b[36m0.8519\u001b[0m        \u001b[32m0.4347\u001b[0m       \u001b[35m0.8448\u001b[0m        \u001b[31m0.4533\u001b[0m  2.1635\n",
      "      6       \u001b[36m0.8540\u001b[0m        \u001b[32m0.4297\u001b[0m       \u001b[35m0.8461\u001b[0m        \u001b[31m0.4516\u001b[0m  2.1888\n",
      "      7       \u001b[36m0.8553\u001b[0m        \u001b[32m0.4258\u001b[0m       \u001b[35m0.8465\u001b[0m        \u001b[31m0.4509\u001b[0m  2.1393\n",
      "      8       \u001b[36m0.8570\u001b[0m        \u001b[32m0.4227\u001b[0m       0.8465        \u001b[31m0.4506\u001b[0m  2.1642\n",
      "      9       \u001b[36m0.8578\u001b[0m        \u001b[32m0.4201\u001b[0m       \u001b[35m0.8468\u001b[0m        0.4507  2.2154\n",
      "     10       \u001b[36m0.8590\u001b[0m        \u001b[32m0.4179\u001b[0m       0.8466        0.4510  2.3116\n",
      "     11       \u001b[36m0.8602\u001b[0m        \u001b[32m0.4160\u001b[0m       \u001b[35m0.8477\u001b[0m        0.4515  2.1901\n",
      "     12       \u001b[36m0.8608\u001b[0m        \u001b[32m0.4143\u001b[0m       0.8474        0.4520  2.1735\n",
      "     13       \u001b[36m0.8618\u001b[0m        \u001b[32m0.4128\u001b[0m       0.8473        0.4526  2.2863\n",
      "     14       \u001b[36m0.8626\u001b[0m        \u001b[32m0.4115\u001b[0m       0.8472        0.4532  2.2163\n",
      "     15       \u001b[36m0.8631\u001b[0m        \u001b[32m0.4103\u001b[0m       0.8471        0.4539  2.0270\n",
      "Test accuracy: 0.8363 \n",
      "\n",
      "### FFNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.7826\u001b[0m        \u001b[32m0.5931\u001b[0m       \u001b[35m0.8327\u001b[0m        \u001b[31m0.4488\u001b[0m  3.3172\n",
      "      2       \u001b[36m0.8381\u001b[0m        \u001b[32m0.4572\u001b[0m       \u001b[35m0.8459\u001b[0m        \u001b[31m0.4277\u001b[0m  3.8799\n",
      "      3       \u001b[36m0.8500\u001b[0m        \u001b[32m0.4296\u001b[0m       \u001b[35m0.8555\u001b[0m        \u001b[31m0.4243\u001b[0m  3.7612\n",
      "      4       \u001b[36m0.8552\u001b[0m        \u001b[32m0.4174\u001b[0m       \u001b[35m0.8672\u001b[0m        \u001b[31m0.3989\u001b[0m  3.5930\n",
      "      5       \u001b[36m0.8592\u001b[0m        \u001b[32m0.4071\u001b[0m       \u001b[35m0.8672\u001b[0m        0.4075  4.0038\n",
      "      6       \u001b[36m0.8645\u001b[0m        \u001b[32m0.3990\u001b[0m       \u001b[35m0.8702\u001b[0m        \u001b[31m0.3788\u001b[0m  3.7168\n",
      "      7       \u001b[36m0.8661\u001b[0m        \u001b[32m0.3925\u001b[0m       0.8685        0.4108  3.8341\n",
      "      8       \u001b[36m0.8676\u001b[0m        \u001b[32m0.3916\u001b[0m       \u001b[35m0.8740\u001b[0m        0.3865  3.7558\n",
      "      9       \u001b[36m0.8683\u001b[0m        \u001b[32m0.3881\u001b[0m       \u001b[35m0.8752\u001b[0m        0.4069  3.1817\n",
      "     10       \u001b[36m0.8708\u001b[0m        \u001b[32m0.3837\u001b[0m       0.8736        0.4105  3.4485\n",
      "     11       \u001b[36m0.8734\u001b[0m        0.3840       \u001b[35m0.8818\u001b[0m        0.3997  3.6171\n",
      "     12       \u001b[36m0.8739\u001b[0m        \u001b[32m0.3811\u001b[0m       0.8782        0.4131  3.9902\n",
      "     13       \u001b[36m0.8744\u001b[0m        \u001b[32m0.3760\u001b[0m       0.8775        0.4049  3.5877\n",
      "     14       \u001b[36m0.8776\u001b[0m        \u001b[32m0.3731\u001b[0m       0.8742        0.4266  3.2744\n",
      "     15       0.8763        0.3737       0.8784        0.4210  3.9802\n",
      "Test accuracy: 0.8713 \n",
      "\n",
      "### CNN\n",
      "  epoch    train_acc    train_loss    valid_acc    valid_loss     dur\n",
      "-------  -----------  ------------  -----------  ------------  ------\n",
      "      1       \u001b[36m0.8079\u001b[0m        \u001b[32m0.5324\u001b[0m       \u001b[35m0.8864\u001b[0m        \u001b[31m0.3138\u001b[0m  6.8582\n",
      "      2       \u001b[36m0.8775\u001b[0m        \u001b[32m0.3512\u001b[0m       \u001b[35m0.9021\u001b[0m        \u001b[31m0.2723\u001b[0m  7.1190\n",
      "      3       \u001b[36m0.8934\u001b[0m        \u001b[32m0.3089\u001b[0m       \u001b[35m0.9119\u001b[0m        \u001b[31m0.2423\u001b[0m  6.8647\n",
      "      4       \u001b[36m0.9007\u001b[0m        \u001b[32m0.2865\u001b[0m       0.9047        0.2611  7.2493\n",
      "      5       \u001b[36m0.9072\u001b[0m        \u001b[32m0.2727\u001b[0m       0.9089        0.2543  7.2193\n",
      "      6       \u001b[36m0.9120\u001b[0m        \u001b[32m0.2592\u001b[0m       \u001b[35m0.9185\u001b[0m        \u001b[31m0.2338\u001b[0m  7.1084\n",
      "      7       \u001b[36m0.9145\u001b[0m        \u001b[32m0.2519\u001b[0m       0.9068        0.2671  7.1128\n",
      "      8       \u001b[36m0.9181\u001b[0m        \u001b[32m0.2409\u001b[0m       \u001b[35m0.9202\u001b[0m        0.2636  7.1794\n",
      "      9       \u001b[36m0.9193\u001b[0m        \u001b[32m0.2408\u001b[0m       0.9157        0.2507  7.1312\n",
      "     10       \u001b[36m0.9218\u001b[0m        \u001b[32m0.2326\u001b[0m       0.9188        0.2371  7.2552\n",
      "     11       \u001b[36m0.9250\u001b[0m        \u001b[32m0.2257\u001b[0m       0.9190        \u001b[31m0.2338\u001b[0m  7.1150\n",
      "     12       \u001b[36m0.9260\u001b[0m        0.2262       \u001b[35m0.9225\u001b[0m        0.2382  7.1878\n",
      "     13       \u001b[36m0.9270\u001b[0m        \u001b[32m0.2211\u001b[0m       0.9193        0.2566  7.1713\n",
      "     14       \u001b[36m0.9275\u001b[0m        \u001b[32m0.2185\u001b[0m       \u001b[35m0.9236\u001b[0m        0.2352  6.8393\n",
      "     15       \u001b[36m0.9292\u001b[0m        \u001b[32m0.2171\u001b[0m       0.9223        0.2606  6.9869\n",
      "Test accuracy: 0.9173 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# train with more epochs\n",
    "for net_module in [LinNN, FFNN, CNN]:\n",
    "    if net_module == CNN:\n",
    "        net = eval_net(net_module, X_train.reshape(-1, 1, 28, 28), y_train, X_test.reshape(-1, 1, 28, 28), y_test, max_epochs=15)\n",
    "    else:\n",
    "        net = eval_net(net_module, X_train, y_train, X_test, y_test, max_epochs=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
