{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd22c2e6",
   "metadata": {},
   "source": [
    "\n",
    "# Chapter 10\n",
    "\n",
    "# Lab: Deep Learning\n",
    "In this section we  demonstrate how to fit the examples discussed\n",
    "in the text. We use the `Python`  `torch` package, along with the\n",
    "`pytorch_lightning` package which provides utilities to simplify\n",
    "fitting and evaluating models.  This code can be impressively fast\n",
    "with certain  special processors,  such as Apple’s new M1 chip. The package is well-structured, flexible, and will feel comfortable\n",
    "to `Python`  users. A good companion is the site\n",
    "[pytorch.org/tutorials](https://pytorch.org/tutorials/beginner/basics/intro.html).\n",
    "Much of our code is adapted from there, as well as the `pytorch_lightning` documentation. {The precise URLs at the time of writing are <https://pytorch.org/tutorials/beginner/basics/intro.html> and <https://pytorch-lightning.readthedocs.io/en/latest/>.}\n",
    "\n",
    "We start with several standard imports that we have  seen in other contexts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aeb913e2",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np, pandas as pd\n",
    "from matplotlib.pyplot import subplots\n",
    "from sklearn.linear_model import \\\n",
    "     (LinearRegression,\n",
    "      LogisticRegression,\n",
    "      Lasso)\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.pipeline import Pipeline\n",
    "from ISLP import load_data\n",
    "from ISLP.models import ModelSpec as MS\n",
    "from sklearn.model_selection import \\\n",
    "     (train_test_split,\n",
    "      GridSearchCV)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a4ac8d",
   "metadata": {},
   "source": [
    "### Torch-Specific Imports\n",
    "There are a number of imports for `torch`. (These are not\n",
    "included with `ISLP`, so must be installed separately.)\n",
    "First we import the main library\n",
    "and essential tools used to specify sequentially-structured networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3a3e91e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.optim import RMSprop\n",
    "from torch.utils.data import TensorDataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e89c676",
   "metadata": {},
   "source": [
    "There are several other helper packages for `torch`. For instance,\n",
    "the `torchmetrics` package has utilities to compute\n",
    "various metrics to evaluate performance when fitting\n",
    "a model. The `torchinfo` package provides a useful\n",
    "summary of the layers of a model. We use the `read_image()`\n",
    "function when loading test images in Section 10.9.4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3f4fdb8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchmetrics import (MeanAbsoluteError,\n",
    "                          R2Score)\n",
    "from torchinfo import summary\n",
    "from torchvision.io import read_image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b12e5c",
   "metadata": {},
   "source": [
    "The package `pytorch_lightning` is a somewhat higher-level\n",
    "interface to `torch` that simplifies the specification and\n",
    "fitting of\n",
    "models by reducing the amount of boilerplate code needed\n",
    "(compared to using `torch` alone)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "71cc6f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.loggers import CSVLogger\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8394fc0a",
   "metadata": {},
   "source": [
    "In order to reproduce results we use `seed_everything()`. We will also instruct `torch` to use deterministic algorithms\n",
    "where possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd1290f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 0\n"
     ]
    }
   ],
   "source": [
    "from pytorch_lightning.utilities.seed import seed_everything\n",
    "seed_everything(0, workers=True)\n",
    "torch.use_deterministic_algorithms(True, warn_only=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f29a6c10",
   "metadata": {},
   "source": [
    "We will use several datasets shipped with `torchvision` for our\n",
    "examples: a pretrained network for image classification,\n",
    "as well as some transforms used for preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c85308d1",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "from torchvision.datasets import MNIST, CIFAR100\n",
    "from torchvision.models import (resnet50,\n",
    "                                ResNet50_Weights)\n",
    "from torchvision.transforms import (Resize,\n",
    "                                    Normalize,\n",
    "                                    CenterCrop,\n",
    "                                    ToTensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ca24f46",
   "metadata": {},
   "source": [
    "We have provided a few utilities in `ISLP` specifically for this lab.\n",
    "The `SimpleDataModule` and `SimpleModule` are simple\n",
    "versions of objects used in `pytorch_lightning`, the\n",
    "high-level module for fitting `torch` models. Although more advanced\n",
    "uses such as computing on graphical processing units (GPUs) and parallel data processing\n",
    "are possible in this module, we will not be focusing much on these\n",
    "in this lab. The `ErrorTracker` handles\n",
    "collections of targets and predictions over each mini-batch\n",
    "in the validation or test stage, allowing computation\n",
    "of the metric over the entire validation or test data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4de9f03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ISLP.torch import (SimpleDataModule,\n",
    "                        SimpleModule,\n",
    "                        ErrorTracker,\n",
    "                        rec_num_workers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d81a775a",
   "metadata": {},
   "source": [
    "In addition we have included some helper\n",
    "functions to load the\n",
    "`IMDb` database, as well as a lookup that maps integers\n",
    "to particular keys in the database. We’ve included\n",
    "a slightly modified copy of the preprocessed\n",
    "`IMDb` data from `keras`, a separate package\n",
    "for fitting deep learning models. This saves us significant\n",
    "preprocessing and allows us to focus on specifying and fitting\n",
    "the models themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ef95fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ISLP.torch.imdb import (load_lookup,\n",
    "                             load_tensor,\n",
    "                             load_sparse,\n",
    "                             load_sequential)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d760ad",
   "metadata": {},
   "source": [
    "Finally, we introduce some utility imports  not directly related to\n",
    "`torch`.\n",
    "The `glob()` function from the `glob` module is used\n",
    "to find all files matching wildcard characters, which we will use\n",
    "in our example applying the `ResNet50` model\n",
    "to some of our own images.\n",
    "The `json` module will be used to load\n",
    "a JSON file for looking up classes to identify the labels of the\n",
    "pictures in the `ResNet50` example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "493821a9",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import json\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c415e4e5",
   "metadata": {},
   "source": [
    "## Single Layer Network on Hitters Data\n",
    "We start by fitting the models in Section 10.6 on the `Hitters` data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a427e224",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "Hitters = load_data('Hitters').dropna()\n",
    "n = Hitters.shape[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1ca2542",
   "metadata": {},
   "source": [
    " We will fit two linear models (least squares  and lasso) and  compare their performance\n",
    "to that of a neural network. For this comparison we will use mean absolute error on a validation dataset.\n",
    "\\begin{equation*}\n",
    "\\begin{split}\n",
    "\\mbox{MAE}(y,\\hat{y}) = \\frac{1}{n} \\sum_{i=1}^n |y_i-\\hat{y}_i|.\n",
    "\\end{split}\n",
    "\\end{equation*}\n",
    "We set up the model matrix and the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "feac4b15",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "model = MS(Hitters.columns.drop('Salary'), intercept=False)\n",
    "X = model.fit_transform(Hitters).to_numpy()\n",
    "Y = Hitters['Salary'].to_numpy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad0d90b5",
   "metadata": {},
   "source": [
    "The `to_numpy()`  method above converts `pandas`\n",
    "data frames or series to `numpy` arrays.\n",
    "We do this because we will need to  use `sklearn` to fit the lasso model,\n",
    "and it requires this conversion. \n",
    "We also use  a linear regression method from `sklearn`, rather than the method\n",
    "in Chapter~3 from `statsmodels`, to facilitate the comparisons."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35acacc2",
   "metadata": {},
   "source": [
    "We now split the data into test and training, fixing the random\n",
    "state used by `sklearn` to do the split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6a75a640",
   "metadata": {},
   "outputs": [],
   "source": [
    "(X_train, \n",
    " X_test,\n",
    " Y_train,\n",
    " Y_test) = train_test_split(X,\n",
    "                            Y,\n",
    "                            test_size=1/3,\n",
    "                            random_state=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc24bfb",
   "metadata": {},
   "source": [
    "### Linear Models\n",
    "We fit the linear model and evaluate the test error directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "32d8de64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "259.71528833146243"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_lm = LinearRegression().fit(X_train, Y_train)\n",
    "Yhat_test = hit_lm.predict(X_test)\n",
    "np.abs(Yhat_test - Y_test).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d60340ef",
   "metadata": {},
   "source": [
    "Next we fit the lasso using `sklearn`. We are using\n",
    "mean absolute error to select and evaluate a model, rather than mean squared error.\n",
    "The specialized solver we used in Section 6.5.2 uses only mean squared error. So here, with a bit more work,  we create a cross-validation grid and perform the cross-validation directly.  \n",
    "\n",
    "We encode a pipeline with two steps: we first normalize the features using a `StandardScaler()` transform,\n",
    "and then fit the lasso without further normalization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "46786db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = StandardScaler(with_mean=True, with_std=True)\n",
    "lasso = Lasso(warm_start=True, max_iter=30000)\n",
    "standard_lasso = Pipeline(steps=[('scaler', scaler),\n",
    "                                 ('lasso', lasso)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ff854a",
   "metadata": {},
   "source": [
    "We need to create a grid of values for $\\lambda$. As is common practice, \n",
    "we choose a grid of 100 values of $\\lambda$, uniform on the log scale from `lam_max` down to  `0.01*lam_max`. Here  `lam_max` is the smallest value of\n",
    "$\\lambda$ with an  all-zero solution. This value equals the largest absolute inner-product between any predictor and the (centered) response. {The derivation of this result is beyond the scope of this book.}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2b55b38b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "X_s = scaler.fit_transform(X_train)\n",
    "n = X_s.shape[0]\n",
    "lam_max = np.fabs(X_s.T.dot(Y_train - Y_train.mean())).max() / n\n",
    "param_grid = {'alpha': np.exp(np.linspace(0, np.log(0.01), 100))\n",
    "             * lam_max}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "977dea00",
   "metadata": {},
   "source": [
    "Note that we had to transform the data first, since the scale of the variables impacts the choice of $\\lambda$.\n",
    "We now perform cross-validation using this sequence of $\\lambda$ values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cc44c559",
   "metadata": {},
   "outputs": [],
   "source": [
    "cv = KFold(10,\n",
    "           shuffle=True,\n",
    "           random_state=1)\n",
    "grid = GridSearchCV(lasso,\n",
    "                    param_grid,\n",
    "                    cv=cv,\n",
    "                    scoring='neg_mean_absolute_error')\n",
    "grid.fit(X_train, Y_train);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca88ebb",
   "metadata": {},
   "source": [
    "We extract the lasso model with best cross-validated mean absolute error, and evaluate its\n",
    "performance on `X_test` and `Y_test`, which were not used in\n",
    "cross-validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1bd1fd13",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "257.2382010799497"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trained_lasso = grid.best_estimator_\n",
    "Yhat_test = trained_lasso.predict(X_test)\n",
    "np.fabs(Yhat_test - Y_test).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f19cd1f2",
   "metadata": {},
   "source": [
    "This is similar to the results we got for the linear model fit by least squares. However, these results can vary a lot for different train/test splits; we encourage the reader to try a different seed in code block 12 and rerun the subsequent code up to this point.\n",
    "\n",
    "### Specifying a Network: Classes and Inheritance\n",
    "To fit the neural network, we first set up a model structure\n",
    "that describes the network.\n",
    "Doing so requires us to define new classes specific to the model we wish to fit.\n",
    "Typically this is done in  `pytorch` by sub-classing a generic\n",
    "representation of a network, which is the approach we take here.\n",
    "Although this example is simple, we will go through the steps in some detail, since it will serve us well\n",
    "for the more complex examples to follow.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4b2ccc5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HittersModel(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size):\n",
    "        super(HittersModel, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.sequential = nn.Sequential(\n",
    "            nn.Linear(input_size, 50),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.Linear(50, 1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        return torch.flatten(self.sequential(x))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da217b2",
   "metadata": {},
   "source": [
    "The `class` statement identifies the code chunk as a\n",
    "declaration for a class `HittersModel`\n",
    "that inherits from the  base class `nn.Module`. This base\n",
    "class is ubiquitous in `torch` and represents the\n",
    "mappings in the neural networks.\n",
    "\n",
    "Indented beneath the `class` statement are the methods of this class:\n",
    "in this case `__init__` and `forward`.  The `__init__` method is\n",
    "called when an instance of the class is created as in the cell\n",
    "below. In the methods, `self` always refers to an instance of the\n",
    "class. In the `__init__` method, we have attached two objects to\n",
    "`self` as attributes: `flatten` and `sequential`. These are used in\n",
    "the `forward` method to describe the map that this module implements.\n",
    "\n",
    "There is one additional line in the `__init__` method, which\n",
    "is a call to\n",
    "`super()`. This function allows subclasses (i.e. `HittersModel`)\n",
    "to access methods of the class they inherit from. For example,\n",
    "the class `nn.Module` has its own `__init__` method, which is different from\n",
    "the `HittersModel.__init__()` method we’ve written above.\n",
    "Using `super()` allows us to call the method of the base class. For\n",
    "`torch` models, we will always be making this `super()` call as it is necessary\n",
    "for the model to be properly interpreted by `torch`.\n",
    "\n",
    "The object `nn.Module` has more methods than simply `__init__` and `forward`. These\n",
    "methods are directly accessible to `HittersModel` instances because of this inheritance.\n",
    "One such method we will see shortly is the `eval()` method, used\n",
    "to disable dropout for when we want to evaluate the model on test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "108ad1f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_model = HittersModel(X.shape[1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d750dd",
   "metadata": {},
   "source": [
    "The object `self.sequential` is a composition of four maps. The\n",
    "first maps the 19 features of `Hitters` to 50 dimensions, introducing $50\\times 19+50$ parameters\n",
    "for the weights and *intercept*  of the map (often called the *bias*). This layer\n",
    "is then mapped to a ReLU layer followed by a 40% dropout layer, and finally a\n",
    "linear map down to 1 dimension, again with a bias. The total number of\n",
    "trainable parameters is therefore $50\\times 19+50+50+1=1051$.\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0c250c5",
   "metadata": {},
   "source": [
    "The package `torchinfo` provides a `summary()` function that neatly summarizes\n",
    "this information. We specify the size of the input and see the size\n",
    "of each tensor as it passes through layers of the network. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "2f20059e",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "HittersModel                             [175, 19]                 [175]                     --\n",
       "├─Flatten: 1-1                           [175, 19]                 [175, 19]                 --\n",
       "├─Sequential: 1-2                        [175, 19]                 [175, 1]                  --\n",
       "│    └─Linear: 2-1                       [175, 19]                 [175, 50]                 1,000\n",
       "│    └─ReLU: 2-2                         [175, 50]                 [175, 50]                 --\n",
       "│    └─Dropout: 2-3                      [175, 50]                 [175, 50]                 --\n",
       "│    └─Linear: 2-4                       [175, 50]                 [175, 1]                  51\n",
       "===================================================================================================================\n",
       "Total params: 1,051\n",
       "Trainable params: 1,051\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 0.18\n",
       "===================================================================================================================\n",
       "Input size (MB): 0.01\n",
       "Forward/backward pass size (MB): 0.07\n",
       "Params size (MB): 0.00\n",
       "Estimated Total Size (MB): 0.09\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(hit_model, \n",
    "        input_size=X_train.shape,\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a0eb81",
   "metadata": {},
   "source": [
    "We have truncated the end of the output slightly, here and in subsequent uses.\n",
    "\n",
    "We now need to transform our training data into a form accessible to `torch`.\n",
    "The basic\n",
    "datatype in `torch` is a `tensor`, which is very similar\n",
    "to an `ndarray` from early chapters.\n",
    "We also note here that `torch` typically\n",
    "works with 32-bit (*single precision*)\n",
    "rather than 64-bit (*double precision*) floating point numbers.\n",
    "We therefore convert our data to `np.float32` before\n",
    "forming the tensor.\n",
    "The $X$ and $Y$ tensors are then arranged into a `Dataset`\n",
    "recognized by `torch`\n",
    "using `TensorDataset()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6ca3030d",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "X_train_t = torch.tensor(X_train.astype(np.float32))\n",
    "Y_train_t = torch.tensor(Y_train.astype(np.float32))\n",
    "hit_train = TensorDataset(X_train_t, Y_train_t)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7924f53d",
   "metadata": {},
   "source": [
    "We do the same for the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "86723b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test_t = torch.tensor(X_test.astype(np.float32))\n",
    "Y_test_t = torch.tensor(Y_test.astype(np.float32))\n",
    "hit_test = TensorDataset(X_test_t, Y_test_t)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d055846",
   "metadata": {},
   "source": [
    "Finally, this dataset is passed to a `DataLoader()` which ultimately\n",
    "passes data into our network. While this may seem\n",
    "like a lot of overhead, this structure is helpful for more\n",
    "complex tasks where data may live on different machines,\n",
    "or where data must be passed to a GPU.\n",
    "We provide a helper function `SimpleDataModule()` in `ISLP` to make this task easier for\n",
    "standard usage.\n",
    "One of its arguments is `num_workers`, which indicates\n",
    "how many processes we will use\n",
    "for loading the data. For small\n",
    "data like `Hitters` this will have little effect, but\n",
    "it does provide an advantage for the `MNIST`  and `CIFAR100` examples below.\n",
    "The `torch` package will inspect the process running and determine a\n",
    "maximum number of workers. {This depends on the computing hardware and the number of cores available.} We’ve included a function\n",
    "`rec_num_workers()` to compute this so we know how many\n",
    "workers might be reasonable (here the max was 16)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "999279fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_num_workers = rec_num_workers()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fe60b31",
   "metadata": {},
   "source": [
    "The general training setup in `pytorch_lightning` involves\n",
    "training, validation and test data. These are each\n",
    "represented by different data loaders. During each epoch,\n",
    "we run a training step to learn the model and a validation\n",
    "step to track the error. The test data is typically\n",
    "used at the end of training to evaluate the model.\n",
    "\n",
    "In this case, as we had split only into test and training,\n",
    "we’ll use the test data as validation data with the\n",
    "argument `validation=hit_test`. The\n",
    "`validation` argument can be a float between 0 and 1, an\n",
    "integer, or a\n",
    "`Dataset`. If a float (respectively, integer), it is interpreted\n",
    "as a percentage (respectively number) of the *training* observations to be used for validation.\n",
    "If it is a `Dataset`, it is passed directly to a data loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7bd9cc5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_dm = SimpleDataModule(hit_train,\n",
    "                          hit_test,\n",
    "                          batch_size=32,\n",
    "                          num_workers=min(4, max_num_workers),\n",
    "                          validation=hit_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be1f578",
   "metadata": {},
   "source": [
    "Next we must provide a `pytorch_lightning` module that controls\n",
    "the steps performed during the training process. We provide methods for our\n",
    "`SimpleModule()` that simply record the value\n",
    "of the loss function and any additional\n",
    "metrics at the end of each epoch. These operations\n",
    "are controlled by the methods `SimpleModule.[training/test/validation]_step()`, though\n",
    "we will not be modifying these in our examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5be2f822",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_module = SimpleModule.regression(hit_model,\n",
    "                           metrics={'mae':MeanAbsoluteError()})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e707d8",
   "metadata": {},
   "source": [
    " By using the `SimpleModule.regression()` method,  we indicate that we will use squared-error loss as in\n",
    "(10.23).\n",
    "We have also asked for mean absolute error to be tracked as well\n",
    "in the metrics that are logged.\n",
    "\n",
    "We log our results via `CSVLogger()`, which in this case stores the results in a CSV file within a directory `logs/hitters`. After the fitting is complete, this allows us to load the\n",
    "results as a `pd.DataFrame()` and visualize them below. There are\n",
    "several ways to log the results within `pytorch_lightning`, though\n",
    "we will not cover those here in detail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "87334d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_logger = CSVLogger('logs', name='hitters')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07663639",
   "metadata": {},
   "source": [
    "Finally we are ready to train our model and log the results. We\n",
    "use the `Trainer()` object from `pytorch_lightning`\n",
    "to do this work. The argument `datamodule=hit_dm` tells the trainer\n",
    "how training/validation/test logs are produced,\n",
    "while the first argument `hit_module`\n",
    "specifies the network architecture\n",
    "as well as the training/validation/test steps.\n",
    "The `callbacks` argument allows for\n",
    "several tasks to be carried out at various\n",
    "points while training a model. Here\n",
    "our `ErrorTracker()` callback will enable\n",
    "us to compute validation error while training\n",
    "and, finally, the test error.\n",
    "We now fit the model for 50 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1a474999",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type         | Params\n",
      "---------------------------------------\n",
      "0 | model | HittersModel | 1.1 K \n",
      "1 | loss  | MSELoss      | 0     \n",
      "---------------------------------------\n",
      "1.1 K     Trainable params\n",
      "0         Non-trainable params\n",
      "1.1 K     Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f1f820ee6fa4d81a3736b68611f3bac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hit_trainer = Trainer(deterministic=True,\n",
    "                      max_epochs=50,\n",
    "                      log_every_n_steps=5,\n",
    "                      logger=hit_logger,\n",
    "                      callbacks=[ErrorTracker()])\n",
    "hit_trainer.fit(hit_module, datamodule=hit_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e154e47",
   "metadata": {},
   "source": [
    "At each step of SGD, the algorithm randomly selects 32 training observations for\n",
    "the computation of the gradient. Recall from Section 10.7\n",
    "that an epoch amounts to the number of SGD steps required to process $n$\n",
    "observations. Since the training set has\n",
    "$n=175$, and we specified a `batch_size` of 32 in the construction of  `hit_dm`, an epoch is $175/32=5.5$ SGD steps.\n",
    "\n",
    "After having fit the model, we can evaluate performance on our test\n",
    "data using the `test()` method of our trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "e3b24643",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce9e57deabef4f55aba22efb507872f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "        test_loss             103304.8515625\n",
      "        test_mae            224.26962280273438\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 103304.8515625, 'test_mae': 224.26962280273438}]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_trainer.test(hit_module, datamodule=hit_dm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0220713b",
   "metadata": {},
   "source": [
    "The results of the fit have been logged into a CSV file. We can find the\n",
    "results specific to this run in the `experiment.metrics_file_path`\n",
    "attribute of our logger. Note that each time the model is fit, the logger will output\n",
    "results into a new subdirectory of our directory `logs/hitters`.\n",
    "\n",
    "We now create a plot of the MAE (mean absolute error) as a function of\n",
    "the number of epochs.\n",
    "First we retrieve the logged summaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f9b266e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_results = pd.read_csv(hit_logger.experiment.metrics_file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7ff26d7",
   "metadata": {},
   "source": [
    "Since we will produce similar plots in later examples, we write a\n",
    "simple generic function to produce this plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c5752a0b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "def summary_plot(results,\n",
    "                 ax,\n",
    "                 col='loss',\n",
    "                 valid_legend='Validation',\n",
    "                 training_legend='Training',\n",
    "                 ylabel='Loss',\n",
    "                 fontsize=20):\n",
    "    for (column,\n",
    "         color,\n",
    "         label) in zip([f'train_{col}_epoch',\n",
    "                        f'valid_{col}'],\n",
    "                       ['black',\n",
    "                        'red'],\n",
    "                       [training_legend,\n",
    "                        valid_legend]):\n",
    "        results.plot(x='epoch',\n",
    "                     y=column,\n",
    "                     label=label,\n",
    "                     marker='o',\n",
    "                     color=color,\n",
    "                     ax=ax)\n",
    "    ax.set_xlabel('Epoch')\n",
    "    ax.set_ylabel(ylabel)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "874c2611",
   "metadata": {},
   "source": [
    "We now set up our axes, and use our function to produce the MAE plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ff0c9fa0",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29688 characters]",
      "text/plain": [
       "<Figure size 600x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = subplots(1, 1, figsize=(6, 6))\n",
    "ax = summary_plot(hit_results,\n",
    "                  ax,\n",
    "                  col='mae',\n",
    "                  ylabel='MAE',\n",
    "                  valid_legend='Validation (=Test)')\n",
    "ax.set_ylim([0, 400])\n",
    "ax.set_xticks(np.linspace(0, 50, 11).astype(int));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212c64a3",
   "metadata": {},
   "source": [
    "We can predict directly from the final model, and\n",
    "evaluate its performance on the test data.\n",
    "Before fitting, we call the `eval()` method\n",
    "of `hit_model`.\n",
    "This tells\n",
    "`torch` to effectively consider this model to be fitted, so that\n",
    "we can use it to predict on new data. For our model here,\n",
    "the biggest change is that the dropout layers will\n",
    "be turned off, i.e. no weights will be randomly\n",
    "dropped in predicting on new data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1ee2b61b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(224.2696, grad_fn=<MeanBackward0>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_model.eval() \n",
    "preds = hit_module(X_test_t)\n",
    "torch.abs(Y_test_t - preds).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8140eb0d",
   "metadata": {},
   "source": [
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64fa7609",
   "metadata": {},
   "source": [
    "### Cleanup\n",
    "In setting up our data module, we had initiated\n",
    "several worker processes that will remain running.\n",
    "We delete all references to the torch objects to ensure these processes\n",
    "will be killed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0ca069a3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "del(Hitters,\n",
    "    hit_model, hit_dm,\n",
    "    hit_logger,\n",
    "    hit_test, hit_train,\n",
    "    X, Y,\n",
    "    X_test, X_train,\n",
    "    Y_test, Y_train,\n",
    "    X_test_t, Y_test_t,\n",
    "    hit_trainer, hit_module)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac321b7f",
   "metadata": {},
   "source": [
    "## Multilayer Network on the MNIST Digit Data\n",
    "The `torchvision` package comes with a number of example datasets,\n",
    "including the `MNIST`  digit data. Our first step is to retrieve\n",
    "the training and test data sets; the `MNIST()` function within\n",
    "`torchvision.datasets` is provided for this purpose. The\n",
    "data will be downloaded the first time this function is executed, and stored in the directory `data/MNIST`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "af6a0a33",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset MNIST\n",
       "    Number of datapoints: 60000\n",
       "    Root location: data\n",
       "    Split: Train\n",
       "    StandardTransform\n",
       "Transform: ToTensor()"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(mnist_train, \n",
    " mnist_test) = [MNIST(root='data',\n",
    "                      train=train,\n",
    "                      download=True,\n",
    "                      transform=ToTensor())\n",
    "                for train in [True, False]]\n",
    "mnist_train\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c8aa040",
   "metadata": {},
   "source": [
    "There are 60,000 images in the training data and 10,000 in the test\n",
    "data. The images are $28\\times 28$, and stored as a matrix of pixels. We\n",
    "need to transform each one into a vector.\n",
    "\n",
    "Neural networks are somewhat sensitive to the scale of the inputs, much as ridge and\n",
    "lasso regularization are affected by scaling.  Here the inputs are eight-bit\n",
    "grayscale values between 0 and 255, so we rescale to the unit\n",
    "interval. {Note: eight bits means $2^8$, which equals 256. Since the convention\n",
    "is to start at $0$, the possible values  range from $0$ to $255$.}\n",
    "This transformation, along with some reordering\n",
    "of the axes, is performed by the `ToTensor()` transform\n",
    "from the `torchvision.transforms` package.\n",
    "\n",
    "As in our `Hitters` example, we form a data module\n",
    "from the training and test datasets, setting aside 20%\n",
    "of the training images for validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "0df8a0c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_dm = SimpleDataModule(mnist_train,\n",
    "                            mnist_test,\n",
    "                            validation=0.2,\n",
    "                            num_workers=max_num_workers,\n",
    "                            batch_size=256)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc50b18",
   "metadata": {},
   "source": [
    "Let’s take a look at the data that will get fed into our network. We loop through the first few\n",
    "chunks of the test dataset, breaking after 2 batches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0c7dd6ee",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:  torch.Size([256, 1, 28, 28])\n",
      "Y:  torch.Size([256])\n",
      "X:  torch.Size([256, 1, 28, 28])\n",
      "Y:  torch.Size([256])\n"
     ]
    }
   ],
   "source": [
    "for idx, (X_ ,Y_) in enumerate(mnist_dm.train_dataloader()):\n",
    "    print('X: ', X_.shape)\n",
    "    print('Y: ', Y_.shape)\n",
    "    if idx >= 1:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc80f21e",
   "metadata": {},
   "source": [
    "We see that the $X$ for each batch consists of 256 images of size `1x28x28`.\n",
    "Here the `1` indicates a single channel (greyscale). For RGB images such as `CIFAR100` below,\n",
    "we will see that the `1` in the size will be replaced by `3` for the three RGB channels.\n",
    "\n",
    "Now we are ready to specify our neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "63e48c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MNISTModel, self).__init__()\n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(28*28, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.4))\n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Linear(256, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3))\n",
    "        self._forward = nn.Sequential(\n",
    "            self.layer1,\n",
    "            self.layer2,\n",
    "            nn.Linear(128, 10))\n",
    "    def forward(self, x):\n",
    "        return self._forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a5ffef",
   "metadata": {},
   "source": [
    "We see that in the first layer, each `1x28x28` image is flattened, then mapped to\n",
    "256 dimensions where we apply a ReLU activation with 40% dropout.\n",
    "A second layer maps the first layer’s output down to\n",
    "128 dimensions, applying a ReLU activation with 30% dropout. Finally,\n",
    "the 128 dimensions are mapped down to 10, the number of classes in the\n",
    "`MNIST`  data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "92405826",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_model = MNISTModel()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b16bd23a",
   "metadata": {},
   "source": [
    "We can check that the model produces output of expected size based\n",
    "on our existing batch `X_` above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "2c8f7a48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([256, 10])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mnist_model(X_).size()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da0de828",
   "metadata": {},
   "source": [
    "Let’s take a look at the summary of the model. Instead of an `input_size` we can pass\n",
    "a tensor of correct shape. In this case, we pass through the final\n",
    "batched `X_` from above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "9f502df8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "MNISTModel                               [256, 1, 28, 28]          [256, 10]                 --\n",
       "├─Sequential: 1-1                        [256, 1, 28, 28]          [256, 10]                 --\n",
       "│    └─Sequential: 2-1                   [256, 1, 28, 28]          [256, 256]                --\n",
       "│    │    └─Flatten: 3-1                 [256, 1, 28, 28]          [256, 784]                --\n",
       "│    │    └─Linear: 3-2                  [256, 784]                [256, 256]                200,960\n",
       "│    │    └─ReLU: 3-3                    [256, 256]                [256, 256]                --\n",
       "│    │    └─Dropout: 3-4                 [256, 256]                [256, 256]                --\n",
       "│    └─Sequential: 2-2                   [256, 256]                [256, 128]                --\n",
       "│    │    └─Linear: 3-5                  [256, 256]                [256, 128]                32,896\n",
       "│    │    └─ReLU: 3-6                    [256, 128]                [256, 128]                --\n",
       "│    │    └─Dropout: 3-7                 [256, 128]                [256, 128]                --\n",
       "│    └─Linear: 2-3                       [256, 128]                [256, 10]                 1,290\n",
       "===================================================================================================================\n",
       "Total params: 235,146\n",
       "Trainable params: 235,146\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 60.20\n",
       "===================================================================================================================\n",
       "Input size (MB): 0.80\n",
       "Forward/backward pass size (MB): 0.81\n",
       "Params size (MB): 0.94\n",
       "Estimated Total Size (MB): 2.55\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(mnist_model,\n",
    "        input_data=X_,\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f98b70bb",
   "metadata": {},
   "source": [
    "Having set up both  the model and the data module, fitting this model is\n",
    "now almost identical to the `Hitters` example. In contrast to our regression model, here we will use the\n",
    "`SimpleModule.classification()` method which\n",
    "uses the  cross-entropy loss function instead of mean squared error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "d1a7590f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_module = SimpleModule.classification(mnist_model)\n",
    "mnist_logger = CSVLogger('logs', name='MNIST')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca665912",
   "metadata": {},
   "source": [
    "Now we are ready to go. The final step is to supply training data, and fit the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "e2fe6de3",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | MNISTModel       | 235 K \n",
      "1 | loss  | CrossEntropyLoss | 0     \n",
      "-------------------------------------------\n",
      "235 K     Trainable params\n",
      "0         Non-trainable params\n",
      "235 K     Total params\n",
      "0.941     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82f9178e94874e7e9cbb3b27cc1c351f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
     ]
    }
   ],
   "source": [
    "mnist_trainer = Trainer(deterministic=True,\n",
    "                        max_epochs=30,\n",
    "                        logger=mnist_logger,\n",
    "                        callbacks=[ErrorTracker()])\n",
    "mnist_trainer.fit(mnist_module,\n",
    "                  datamodule=mnist_dm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b4a4e0c",
   "metadata": {},
   "source": [
    "We have suppressed the output here, which is a progress report on the\n",
    "fitting of the model, grouped by epoch. This is very useful, since on\n",
    "large datasets fitting can take time. Fitting this model took 245\n",
    "seconds on a MacBook Pro with an Apple M1 Pro chip with 10 cores and 16 GB of RAM.\n",
    "Here we specified a\n",
    "validation split of 20%, so training is actually performed on\n",
    "80% of the 60,000 observations in the training set. This is an\n",
    "alternative to actually supplying validation data, like we did for the `Hitters` data.\n",
    "SGD  uses batches\n",
    "of 256 observations in computing the gradient, and doing the\n",
    "arithmetic, we see that an epoch corresponds to 188 gradient steps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25528d72",
   "metadata": {},
   "source": [
    "`SimpleModule.classification()` includes\n",
    "an accuracy metric by default. Other\n",
    "classification metrics can be added from `torchmetrics`.\n",
    "We will use  our `summary_plot()` function to display \n",
    "accuracy across epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "73755987",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22272 characters]",
      "text/plain": [
       "<Figure size 600x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mnist_results = pd.read_csv(mnist_logger.experiment.metrics_file_path)\n",
    "fig, ax = subplots(1, 1, figsize=(6, 6))\n",
    "summary_plot(mnist_results,\n",
    "             ax,\n",
    "             col='accuracy',\n",
    "             ylabel='Accuracy')\n",
    "ax.set_ylim([0.5, 1])\n",
    "ax.set_ylabel('Accuracy')\n",
    "ax.set_xticks(np.linspace(0, 30, 7).astype(int));\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba63a8da",
   "metadata": {},
   "source": [
    "Once again we evaluate the accuracy using the `test()` method of our trainer. This model achieves\n",
    "97% accuracy on the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "f5269c40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b584fa2b09df4c20939181c016e85bc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "      test_accuracy         0.9681000113487244\n",
      "        test_loss           0.14706705510616302\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.14706705510616302, 'test_accuracy': 0.9681000113487244}]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mnist_trainer.test(mnist_module,\n",
    "                   datamodule=mnist_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bca75042",
   "metadata": {},
   "source": [
    "Table 10.1 also reports the error rates resulting from LDA (Chapter 4) and multiclass logistic\n",
    "regression. For LDA we refer the reader to Section 4.7.3.\n",
    "Although we could use the `sklearn` function `LogisticRegression()` to fit  \n",
    "multiclass logistic regression, we are set up here to fit such a model\n",
    "with `torch`.\n",
    "We just have an input layer and an output layer, and omit the hidden layers!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "97a0b304",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNIST_MLR(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MNIST_MLR, self).__init__()\n",
    "        self.linear = nn.Sequential(nn.Flatten(),\n",
    "                                    nn.Linear(784, 10))\n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "mlr_model = MNIST_MLR()\n",
    "mlr_module = SimpleModule.classification(mlr_model)\n",
    "mlr_logger = CSVLogger('logs', name='MNIST_MLR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "ea685183",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | MNIST_MLR        | 7.9 K \n",
      "1 | loss  | CrossEntropyLoss | 0     \n",
      "-------------------------------------------\n",
      "7.9 K     Trainable params\n",
      "0         Non-trainable params\n",
      "7.9 K     Total params\n",
      "0.031     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02fd79458dc34808bb403f8b8475a62d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c827fce3589462294778b990fc78f2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9fcb28681b34dc6921cad2630ca648f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e2dd237b55d14ad4a871888e4ecde25d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d8dc977673646d9ad09d0c5b7e4903a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
     ]
    }
   ],
   "source": [
    "mlr_trainer = Trainer(deterministic=True,\n",
    "                      max_epochs=30,\n",
    "                      callbacks=[ErrorTracker()])\n",
    "mlr_trainer.fit(mlr_module, datamodule=mnist_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c376e402",
   "metadata": {},
   "source": [
    "We fit the model just as before and compute the test results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c0bd63e3",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92b4035b1656456d88d929aac0d46d71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "      test_accuracy         0.9240999817848206\n",
      "        test_loss           0.3186827003955841\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.3186827003955841, 'test_accuracy': 0.9240999817848206}]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlr_trainer.test(mlr_module,\n",
    "                 datamodule=mnist_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e99a86",
   "metadata": {},
   "source": [
    "The accuracy is above 90% even for this pretty simple model.\n",
    "\n",
    "As in the `Hitters` example, we delete some of\n",
    "the objects we created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "6b0d3daa",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "del(mnist_test,\n",
    "    mnist_train,\n",
    "    mnist_model,\n",
    "    mnist_dm,\n",
    "    mnist_trainer,\n",
    "    mnist_module,\n",
    "    mnist_results,\n",
    "    mlr_model,\n",
    "    mlr_module,\n",
    "    mlr_trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e370989",
   "metadata": {},
   "source": [
    "## Convolutional Neural Networks\n",
    "In this section we fit a CNN to the `CIFAR100` data, which is available in the `torchvision`\n",
    "package. It is arranged in a similar fashion as the `MNIST` data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "67517b11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "(cifar_train, \n",
    " cifar_test) = [CIFAR100(root=\"data\",\n",
    "                         train=train,\n",
    "                         download=True)\n",
    "             for train in [True, False]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ee7b040e",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = ToTensor()\n",
    "cifar_train_X = torch.stack([transform(x) for x in\n",
    "                            cifar_train.data])\n",
    "cifar_test_X = torch.stack([transform(x) for x in\n",
    "                            cifar_test.data])\n",
    "cifar_train = TensorDataset(cifar_train_X,\n",
    "                            torch.tensor(cifar_train.targets))\n",
    "cifar_test = TensorDataset(cifar_test_X,\n",
    "                            torch.tensor(cifar_test.targets))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bba84d8",
   "metadata": {},
   "source": [
    "The `CIFAR100` dataset consists of 50,000 training images, each represented by a three-dimensional tensor:\n",
    "each three-color image is represented as a set of three channels, each of which consists of\n",
    "$32\\times 32$ eight-bit pixels. We standardize as we did for the\n",
    "digits, but keep the array structure. This is accomplished with the `ToTensor()` transform.\n",
    "\n",
    "Creating the data module is similar to the `MNIST`  example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bd9e74ad",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "cifar_dm = SimpleDataModule(cifar_train,\n",
    "                            cifar_test,\n",
    "                            validation=0.2,\n",
    "                            num_workers=max_num_workers,\n",
    "                            batch_size=128)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee8883b",
   "metadata": {},
   "source": [
    "We again look at the shape of typical batches in our data loaders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "a553b926",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X:  torch.Size([128, 3, 32, 32])\n",
      "Y:  torch.Size([128])\n",
      "X:  torch.Size([128, 3, 32, 32])\n",
      "Y:  torch.Size([128])\n"
     ]
    }
   ],
   "source": [
    "for idx, (X_ ,Y_) in enumerate(cifar_dm.train_dataloader()):\n",
    "    print('X: ', X_.shape)\n",
    "    print('Y: ', Y_.shape)\n",
    "    if idx >= 1:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b25a138",
   "metadata": {},
   "source": [
    "Before we start, we look at some of the training images; similar code produced\n",
    "Figure 10.5 on page  445. The example below also illustrates\n",
    "that `TensorDataset` objects can be indexed with integers --- we are choosing\n",
    "random images from the training data by indexing `cifar_train`. In order to display correctly,\n",
    "we must reorder the dimensions by a call to `np.transpose()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "94885e68",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 143300 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 25 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = subplots(5, 5, figsize=(10,10))\n",
    "rng = np.random.default_rng(4)\n",
    "indices = rng.choice(np.arange(len(cifar_train)), 25,\n",
    "                     replace=False).reshape((5,5))\n",
    "for i in range(5):\n",
    "    for j in range(5):\n",
    "        idx = indices[i,j]\n",
    "        axes[i,j].imshow(np.transpose(cifar_train[idx][0],\n",
    "                                      [1,2,0]),\n",
    "                                      interpolation=None)\n",
    "        axes[i,j].set_xticks([])\n",
    "        axes[i,j].set_yticks([])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04fda6b8",
   "metadata": {},
   "source": [
    "Here the `imshow()` method recognizes from the shape of its argument that it is a 3-dimensional array, with the last dimension indexing the three RGB color channels.\n",
    "\n",
    "We specify a moderately-sized  CNN for\n",
    "demonstration purposes, similar in structure to Figure 10.8.\n",
    "We use several layers, each consisting of  convolution, ReLU, and max-pooling steps.\n",
    "We first define a module that defines one of these layers. As in our\n",
    "previous examples, we overwrite the `__init__()` and `forward()` methods\n",
    "of `nn.Module`. This user-defined  module can now be used in ways just like\n",
    "`nn.Linear()` or `nn.Dropout()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "cdc2528a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BuildingBlock(nn.Module):\n",
    "\n",
    "    def __init__(self,\n",
    "                 in_channels,\n",
    "                 out_channels):\n",
    "\n",
    "        super(BuildingBlock, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels=in_channels,\n",
    "                              out_channels=out_channels,\n",
    "                              kernel_size=(3,3),\n",
    "                              padding='same')\n",
    "        self.activation = nn.ReLU()\n",
    "        self.pool = nn.MaxPool2d(kernel_size=(2,2))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.pool(self.activation(self.conv(x)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31891282",
   "metadata": {},
   "source": [
    "Notice that we used the `padding = \"same\"` argument to\n",
    "`nn.Conv2d()`, which ensures that the output channels have the\n",
    "same dimension as the input channels. There are 32 channels in the first\n",
    "hidden layer, in contrast to the three channels in the input layer. We\n",
    "use a $3\\times 3$ convolution filter for each channel in all the layers. Each\n",
    "convolution is followed by a max-pooling layer over $2\\times2$\n",
    "blocks.\n",
    "\n",
    "In forming our deep learning model for the `CIFAR100` data, we use several of our `BuildingBlock()`\n",
    "modules sequentially. This simple example\n",
    "illustrates some of the power of `torch`. Users can\n",
    "define modules of their own, which can be combined in other\n",
    "modules. Ultimately, everything is fit by a generic trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "845be1ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFARModel(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(CIFARModel, self).__init__()\n",
    "        sizes = [(3,32),\n",
    "                 (32,64),\n",
    "                 (64,128),\n",
    "                 (128,256)]\n",
    "        self.conv = nn.Sequential(*[BuildingBlock(in_, out_)\n",
    "                                    for in_, out_ in sizes])\n",
    "\n",
    "        self.output = nn.Sequential(nn.Dropout(0.5),\n",
    "                                    nn.Linear(2*2*256, 512),\n",
    "                                    nn.ReLU(),\n",
    "                                    nn.Linear(512, 100))\n",
    "    def forward(self, x):\n",
    "        val = self.conv(x)\n",
    "        val = torch.flatten(val, start_dim=1)\n",
    "        return self.output(val)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ed2c036",
   "metadata": {},
   "source": [
    "We  build the model and look at the summary. (We had created examples of `X_` earlier.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "be768cbb",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "CIFARModel                               [128, 3, 32, 32]          [128, 100]                --\n",
       "├─Sequential: 1-1                        [128, 3, 32, 32]          [128, 256, 2, 2]          --\n",
       "│    └─BuildingBlock: 2-1                [128, 3, 32, 32]          [128, 32, 16, 16]         --\n",
       "│    │    └─Conv2d: 3-1                  [128, 3, 32, 32]          [128, 32, 32, 32]         896\n",
       "│    │    └─ReLU: 3-2                    [128, 32, 32, 32]         [128, 32, 32, 32]         --\n",
       "│    │    └─MaxPool2d: 3-3               [128, 32, 32, 32]         [128, 32, 16, 16]         --\n",
       "│    └─BuildingBlock: 2-2                [128, 32, 16, 16]         [128, 64, 8, 8]           --\n",
       "│    │    └─Conv2d: 3-4                  [128, 32, 16, 16]         [128, 64, 16, 16]         18,496\n",
       "│    │    └─ReLU: 3-5                    [128, 64, 16, 16]         [128, 64, 16, 16]         --\n",
       "│    │    └─MaxPool2d: 3-6               [128, 64, 16, 16]         [128, 64, 8, 8]           --\n",
       "│    └─BuildingBlock: 2-3                [128, 64, 8, 8]           [128, 128, 4, 4]          --\n",
       "│    │    └─Conv2d: 3-7                  [128, 64, 8, 8]           [128, 128, 8, 8]          73,856\n",
       "│    │    └─ReLU: 3-8                    [128, 128, 8, 8]          [128, 128, 8, 8]          --\n",
       "│    │    └─MaxPool2d: 3-9               [128, 128, 8, 8]          [128, 128, 4, 4]          --\n",
       "│    └─BuildingBlock: 2-4                [128, 128, 4, 4]          [128, 256, 2, 2]          --\n",
       "│    │    └─Conv2d: 3-10                 [128, 128, 4, 4]          [128, 256, 4, 4]          295,168\n",
       "│    │    └─ReLU: 3-11                   [128, 256, 4, 4]          [128, 256, 4, 4]          --\n",
       "│    │    └─MaxPool2d: 3-12              [128, 256, 4, 4]          [128, 256, 2, 2]          --\n",
       "├─Sequential: 1-2                        [128, 1024]               [128, 100]                --\n",
       "│    └─Dropout: 2-5                      [128, 1024]               [128, 1024]               --\n",
       "│    └─Linear: 2-6                       [128, 1024]               [128, 512]                524,800\n",
       "│    └─ReLU: 2-7                         [128, 512]                [128, 512]                --\n",
       "│    └─Linear: 2-8                       [128, 512]                [128, 100]                51,300\n",
       "===================================================================================================================\n",
       "Total params: 964,516\n",
       "Trainable params: 964,516\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 2.01\n",
       "===================================================================================================================\n",
       "Input size (MB): 1.57\n",
       "Forward/backward pass size (MB): 63.54\n",
       "Params size (MB): 3.86\n",
       "Estimated Total Size (MB): 68.97\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cifar_model = CIFARModel()\n",
    "summary(cifar_model,\n",
    "        input_data=X_,\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6320d294",
   "metadata": {},
   "source": [
    "The total number of trainable parameters is 964,516.\n",
    "By studying the size of the parameters, we can see that the channels halve in both\n",
    "dimensions\n",
    "after each of these max-pooling operations. After the last of these we\n",
    "have a layer with  256 channels of dimension $2\\times 2$. These are then\n",
    "flattened to a dense layer of size 1,024;\n",
    "in other words, each of the $2\\times 2$ matrices is turned into a\n",
    "$4$-vector, and put side-by-side in one layer. This is followed by a\n",
    "dropout regularization layer,  then\n",
    "another dense layer of size 512, and finally, the\n",
    "output layer.\n",
    "\n",
    "Up to now, we have been using a default\n",
    "optimizer in `SimpleModule()`. For these data,\n",
    "experiments show that a smaller learning rate performs\n",
    "better than the default 0.01. We use a\n",
    "custom optimizer here with a learning rate of 0.001.\n",
    "Besides this, the logging and training\n",
    "follow a similar pattern to our previous examples. The optimizer\n",
    "takes an argument `params` that informs\n",
    "the optimizer which parameters are involved in SGD (stochastic gradient descent).\n",
    "\n",
    "We saw earlier that entries of a module’s parameters are tensors. In passing\n",
    "the parameters to the optimizer we are doing more than\n",
    "simply passing arrays; part of the structure of the graph\n",
    "is encoded in the tensors themselves.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "67e3e2d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "cifar_optimizer = RMSprop(cifar_model.parameters(), lr=0.001)\n",
    "cifar_module = SimpleModule.classification(cifar_model,\n",
    "                                    optimizer=cifar_optimizer)\n",
    "cifar_logger = CSVLogger('logs', name='CIFAR100')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "1ea339ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | CIFARModel       | 964 K \n",
      "1 | loss  | CrossEntropyLoss | 0     \n",
      "-------------------------------------------\n",
      "964 K     Trainable params\n",
      "0         Non-trainable params\n",
      "964 K     Total params\n",
      "3.858     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fdb66e85a34e49afafdf51191e247d54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
     ]
    }
   ],
   "source": [
    "cifar_trainer = Trainer(deterministic=True,\n",
    "                        max_epochs=30,\n",
    "                        logger=cifar_logger,\n",
    "                        callbacks=[ErrorTracker()])\n",
    "cifar_trainer.fit(cifar_module,\n",
    "                  datamodule=cifar_dm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "977e27d9",
   "metadata": {},
   "source": [
    "This model takes 10 minutes or more to run and achieves about 42% accuracy on the test\n",
    "data. Although this is not terrible for 100-class data (a random\n",
    "classifier gets 1% accuracy), searching the web we see results around\n",
    "75%. Typically it takes a lot of architecture carpentry,\n",
    "fiddling with regularization, and time, to achieve such results.\n",
    "\n",
    "Let’s take a look at the validation and training accuracy\n",
    "across epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "0a9068d9",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25044 characters]",
      "text/plain": [
       "<Figure size 600x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "log_path = cifar_logger.experiment.metrics_file_path\n",
    "cifar_results = pd.read_csv(log_path)\n",
    "fig, ax = subplots(1, 1, figsize=(6, 6))\n",
    "summary_plot(cifar_results,\n",
    "             ax,\n",
    "             col='accuracy',\n",
    "             ylabel='Accuracy')\n",
    "ax.set_xticks(np.linspace(0, 10, 6).astype(int))\n",
    "ax.set_ylabel('Accuracy')\n",
    "ax.set_ylim([0, 1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c48006a",
   "metadata": {},
   "source": [
    "Finally, we evaluate our model on our test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "e71eff9b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e804034c1964ed38bbe2ca376491f54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "      test_accuracy         0.40880000591278076\n",
      "        test_loss           2.4764862060546875\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 2.4764862060546875, 'test_accuracy': 0.40880000591278076}]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cifar_trainer.test(cifar_module,\n",
    "                   datamodule=cifar_dm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d03b891",
   "metadata": {},
   "source": [
    "### Hardware Acceleration\n",
    "As deep learning has become ubiquitous in machine learning, hardware\n",
    "manufacturers have produced special libraries that can\n",
    "often speed up the gradient-descent steps.\n",
    "\n",
    "For instance, Mac OS devices with the M1 chip may have the *Metal* programming framework\n",
    "enabled, which can speed up the  `torch` \n",
    "computations. We present an example of how to use this acceleration.\n",
    "\n",
    "The main changes are to the `Trainer()` call as well as to the metrics\n",
    "that will be evaluated on the data. These metrics must be told  where\n",
    "the data will be located at evaluation time. This is\n",
    "accomplished with a call to the `to()` method of the metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "58eae430",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | CIFARModel       | 964 K \n",
      "1 | loss  | CrossEntropyLoss | 0     \n",
      "-------------------------------------------\n",
      "964 K     Trainable params\n",
      "0         Non-trainable params\n",
      "964 K     Total params\n",
      "3.858     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ebeb3882e14340e8b43876442c7e36d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchmetrics/utilities/checks.py:49: UserWarning: MPS: no support for int64 min/max ops, casting it to int32 (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/ReduceOps.mm:1271.)\n",
      "  if ignore_index is None and target.min() < 0:\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3436eeffcb544f009e93641c9b125c58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "try:\n",
    "    for name, metric in cifar_module.metrics.items():\n",
    "        cifar_module.metrics[name] = metric.to('mps')\n",
    "    cifar_trainer_mps = Trainer(accelerator='mps',\n",
    "                                deterministic=True,\n",
    "                                max_epochs=30)\n",
    "    cifar_trainer_mps.fit(cifar_module,\n",
    "                          datamodule=cifar_dm)\n",
    "    cifar_trainer_mps.test(cifar_module,\n",
    "                          datamodule=cifar_dm)\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7227ef",
   "metadata": {},
   "source": [
    "This yields approximately two- or three-fold  acceleration for each epoch.\n",
    "We have protected this code block using `try:` and `except:`\n",
    "clauses; if it works, we get the speedup, if it fails, nothing happens."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4328a190",
   "metadata": {},
   "source": [
    "## Using Pretrained CNN Models\n",
    "We now show how to use a CNN pretrained on the  `imagenet` database to classify natural\n",
    "images, and demonstrate how we produced Figure 10.10.\n",
    "We copied six JPEG images from a digital photo album into the\n",
    "directory `book_images`. These images are available\n",
    "from the data section of  <www.statlearning.com>, the ISLP book website. Download `book_images.zip`; when\n",
    "clicked it creates the `book_images` directory. \n",
    "\n",
    "The pretrained network we use is called `resnet50`; specification details can be found on the web.\n",
    "We will read in the images, and\n",
    "convert them into the array format expected by the `torch`\n",
    "software to match the specifications in `resnet50`. \n",
    "The conversion involves a resize, a crop and then a predefined standardization for each of the three channels.\n",
    "We now read in the images and preprocess them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "638a1c8c",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 3, 224, 224])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resize = Resize((232,232))\n",
    "crop = CenterCrop(224)\n",
    "normalize = Normalize([0.485,0.456,0.406],\n",
    "                      [0.229,0.224,0.225])\n",
    "imgfiles = sorted([f for f in glob('book_images/*')])\n",
    "imgs = torch.stack([torch.div(crop(resize(read_image(f))), 255)\n",
    "                    for f in imgfiles])\n",
    "imgs = normalize(imgs)\n",
    "imgs.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c5aec2",
   "metadata": {},
   "source": [
    "We now set up the trained network with the weights we read in code block~6. The model has 50 layers, with a fair bit of complexity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "7776ed1e",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "ResNet                                   [6, 3, 224, 224]          [6, 1000]                 --\n",
       "├─Conv2d: 1-1                            [6, 3, 224, 224]          [6, 64, 112, 112]         9,408\n",
       "├─BatchNorm2d: 1-2                       [6, 64, 112, 112]         [6, 64, 112, 112]         128\n",
       "├─ReLU: 1-3                              [6, 64, 112, 112]         [6, 64, 112, 112]         --\n",
       "├─MaxPool2d: 1-4                         [6, 64, 112, 112]         [6, 64, 56, 56]           --\n",
       "├─Sequential: 1-5                        [6, 64, 56, 56]           [6, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-1                   [6, 64, 56, 56]           [6, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-1                  [6, 64, 56, 56]           [6, 64, 56, 56]           4,096\n",
       "│    │    └─BatchNorm2d: 3-2             [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-3                    [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-4                  [6, 64, 56, 56]           [6, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-5             [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-6                    [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-7                  [6, 64, 56, 56]           [6, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-8             [6, 256, 56, 56]          [6, 256, 56, 56]          512\n",
       "│    │    └─Sequential: 3-9              [6, 64, 56, 56]           [6, 256, 56, 56]          16,896\n",
       "│    │    └─ReLU: 3-10                   [6, 256, 56, 56]          [6, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-2                   [6, 256, 56, 56]          [6, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-11                 [6, 256, 56, 56]          [6, 64, 56, 56]           16,384\n",
       "│    │    └─BatchNorm2d: 3-12            [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-13                   [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-14                 [6, 64, 56, 56]           [6, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-15            [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-16                   [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-17                 [6, 64, 56, 56]           [6, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-18            [6, 256, 56, 56]          [6, 256, 56, 56]          512\n",
       "│    │    └─ReLU: 3-19                   [6, 256, 56, 56]          [6, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-3                   [6, 256, 56, 56]          [6, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-20                 [6, 256, 56, 56]          [6, 64, 56, 56]           16,384\n",
       "│    │    └─BatchNorm2d: 3-21            [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-22                   [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-23                 [6, 64, 56, 56]           [6, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-24            [6, 64, 56, 56]           [6, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-25                   [6, 64, 56, 56]           [6, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-26                 [6, 64, 56, 56]           [6, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-27            [6, 256, 56, 56]          [6, 256, 56, 56]          512\n",
       "│    │    └─ReLU: 3-28                   [6, 256, 56, 56]          [6, 256, 56, 56]          --\n",
       "├─Sequential: 1-6                        [6, 256, 56, 56]          [6, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-4                   [6, 256, 56, 56]          [6, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-29                 [6, 256, 56, 56]          [6, 128, 56, 56]          32,768\n",
       "│    │    └─BatchNorm2d: 3-30            [6, 128, 56, 56]          [6, 128, 56, 56]          256\n",
       "│    │    └─ReLU: 3-31                   [6, 128, 56, 56]          [6, 128, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-32                 [6, 128, 56, 56]          [6, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-33            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-34                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-35                 [6, 128, 28, 28]          [6, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-36            [6, 512, 28, 28]          [6, 512, 28, 28]          1,024\n",
       "│    │    └─Sequential: 3-37             [6, 256, 56, 56]          [6, 512, 28, 28]          132,096\n",
       "│    │    └─ReLU: 3-38                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-5                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-39                 [6, 512, 28, 28]          [6, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-40            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-41                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-42                 [6, 128, 28, 28]          [6, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-43            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-44                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-45                 [6, 128, 28, 28]          [6, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-46            [6, 512, 28, 28]          [6, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-47                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-6                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-48                 [6, 512, 28, 28]          [6, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-49            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-50                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-51                 [6, 128, 28, 28]          [6, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-52            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-53                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-54                 [6, 128, 28, 28]          [6, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-55            [6, 512, 28, 28]          [6, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-56                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-7                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-57                 [6, 512, 28, 28]          [6, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-58            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-59                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-60                 [6, 128, 28, 28]          [6, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-61            [6, 128, 28, 28]          [6, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-62                   [6, 128, 28, 28]          [6, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-63                 [6, 128, 28, 28]          [6, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-64            [6, 512, 28, 28]          [6, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-65                   [6, 512, 28, 28]          [6, 512, 28, 28]          --\n",
       "├─Sequential: 1-7                        [6, 512, 28, 28]          [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-8                   [6, 512, 28, 28]          [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-66                 [6, 512, 28, 28]          [6, 256, 28, 28]          131,072\n",
       "│    │    └─BatchNorm2d: 3-67            [6, 256, 28, 28]          [6, 256, 28, 28]          512\n",
       "│    │    └─ReLU: 3-68                   [6, 256, 28, 28]          [6, 256, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-69                 [6, 256, 28, 28]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-70            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-71                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-72                 [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-73            [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─Sequential: 3-74             [6, 512, 28, 28]          [6, 1024, 14, 14]         526,336\n",
       "│    │    └─ReLU: 3-75                   [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-9                   [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-76                 [6, 1024, 14, 14]         [6, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-77            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-78                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-79                 [6, 256, 14, 14]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-80            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-81                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-82                 [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-83            [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-84                   [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-10                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-85                 [6, 1024, 14, 14]         [6, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-86            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-87                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-88                 [6, 256, 14, 14]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-89            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-90                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-91                 [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-92            [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-93                   [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-11                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-94                 [6, 1024, 14, 14]         [6, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-95            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-96                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-97                 [6, 256, 14, 14]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-98            [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-99                   [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-100                [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-101           [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-102                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-12                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-103                [6, 1024, 14, 14]         [6, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-104           [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-105                  [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-106                [6, 256, 14, 14]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-107           [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-108                  [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-109                [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-110           [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-111                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-13                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-112                [6, 1024, 14, 14]         [6, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-113           [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-114                  [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-115                [6, 256, 14, 14]          [6, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-116           [6, 256, 14, 14]          [6, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-117                  [6, 256, 14, 14]          [6, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-118                [6, 256, 14, 14]          [6, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-119           [6, 1024, 14, 14]         [6, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-120                  [6, 1024, 14, 14]         [6, 1024, 14, 14]         --\n",
       "├─Sequential: 1-8                        [6, 1024, 14, 14]         [6, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-14                  [6, 1024, 14, 14]         [6, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-121                [6, 1024, 14, 14]         [6, 512, 14, 14]          524,288\n",
       "│    │    └─BatchNorm2d: 3-122           [6, 512, 14, 14]          [6, 512, 14, 14]          1,024\n",
       "│    │    └─ReLU: 3-123                  [6, 512, 14, 14]          [6, 512, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-124                [6, 512, 14, 14]          [6, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-125           [6, 512, 7, 7]            [6, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-126                  [6, 512, 7, 7]            [6, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-127                [6, 512, 7, 7]            [6, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-128           [6, 2048, 7, 7]           [6, 2048, 7, 7]           4,096\n",
       "│    │    └─Sequential: 3-129            [6, 1024, 14, 14]         [6, 2048, 7, 7]           2,101,248\n",
       "│    │    └─ReLU: 3-130                  [6, 2048, 7, 7]           [6, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-15                  [6, 2048, 7, 7]           [6, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-131                [6, 2048, 7, 7]           [6, 512, 7, 7]            1,048,576\n",
       "│    │    └─BatchNorm2d: 3-132           [6, 512, 7, 7]            [6, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-133                  [6, 512, 7, 7]            [6, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-134                [6, 512, 7, 7]            [6, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-135           [6, 512, 7, 7]            [6, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-136                  [6, 512, 7, 7]            [6, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-137                [6, 512, 7, 7]            [6, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-138           [6, 2048, 7, 7]           [6, 2048, 7, 7]           4,096\n",
       "│    │    └─ReLU: 3-139                  [6, 2048, 7, 7]           [6, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-16                  [6, 2048, 7, 7]           [6, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-140                [6, 2048, 7, 7]           [6, 512, 7, 7]            1,048,576\n",
       "│    │    └─BatchNorm2d: 3-141           [6, 512, 7, 7]            [6, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-142                  [6, 512, 7, 7]            [6, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-143                [6, 512, 7, 7]            [6, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-144           [6, 512, 7, 7]            [6, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-145                  [6, 512, 7, 7]            [6, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-146                [6, 512, 7, 7]            [6, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-147           [6, 2048, 7, 7]           [6, 2048, 7, 7]           4,096\n",
       "│    │    └─ReLU: 3-148                  [6, 2048, 7, 7]           [6, 2048, 7, 7]           --\n",
       "├─AdaptiveAvgPool2d: 1-9                 [6, 2048, 7, 7]           [6, 2048, 1, 1]           --\n",
       "├─Linear: 1-10                           [6, 2048]                 [6, 1000]                 2,049,000\n",
       "===================================================================================================================\n",
       "Total params: 25,557,032\n",
       "Trainable params: 25,557,032\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 24.54\n",
       "===================================================================================================================\n",
       "Input size (MB): 3.61\n",
       "Forward/backward pass size (MB): 1066.99\n",
       "Params size (MB): 102.23\n",
       "Estimated Total Size (MB): 1172.83\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resnet_model = resnet50(weights=ResNet50_Weights.DEFAULT)\n",
    "summary(resnet_model,\n",
    "        input_data=imgs,\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d29ab45",
   "metadata": {},
   "source": [
    "We set the mode to `eval()` to ensure that the model is ready to predict on new data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "5c8e359d",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=2048, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resnet_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "228a6411",
   "metadata": {},
   "source": [
    "Inspecting the output above, we see that when setting up the\n",
    "`resnet_model`, the authors defined a `Bottleneck`, much like our\n",
    "`BuildingBlock` module.\n",
    "\n",
    "We now feed our six images through the fitted network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "80f19869",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_preds = resnet_model(imgs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7bd653f",
   "metadata": {},
   "source": [
    "Let’s look at the predicted probabilities for each of the top 3 choices. First we compute\n",
    "the probabilities by applying the softmax to the logits in `img_preds`. Note that\n",
    "we have had to call the `detach()` method on the tensor `img_preds` in order to convert\n",
    "it to our a more familiar `ndarray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "6892597d",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_probs = np.exp(np.asarray(img_preds.detach()))\n",
    "img_probs /= img_probs.sum(1)[:,None]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0f07763",
   "metadata": {},
   "source": [
    "In order to see the class labels, we must download the index file associated with `imagenet`. {This is avalable from the book website and  [s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json](https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json).}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "3eb1e1b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "labs = json.load(open('imagenet_class_index.json'))\n",
    "class_labels = pd.DataFrame([(int(k), v[1]) for k, v in \n",
    "                           labs.items()],\n",
    "                           columns=['idx', 'label'])\n",
    "class_labels = class_labels.set_index('idx')\n",
    "class_labels = class_labels.sort_index()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7779e914",
   "metadata": {},
   "source": [
    "We’ll now construct a data frame for each image file\n",
    "with the labels with the three highest probabilities as\n",
    "estimated by the model above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "0e89d251",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image: book_images/flamingo.jpg\n",
      "            label      prob\n",
      "0        flamingo  0.591760\n",
      "1       spoonbill  0.012386\n",
      "2  American_egret  0.002105\n",
      "Image: book_images/hawk.jpg\n",
      "            label      prob\n",
      "0  great_grey_owl  0.287958\n",
      "1            kite  0.039478\n",
      "2        fountain  0.029384\n",
      "Image: book_images/hawk_cropped.jpeg\n",
      "    label      prob\n",
      "0    kite  0.301841\n",
      "1     jay  0.121659\n",
      "2  magpie  0.015511\n",
      "Image: book_images/huey.jpg\n",
      "             label      prob\n",
      "0            Lhasa  0.151153\n",
      "1         Shih-Tzu  0.129804\n",
      "2  Tibetan_terrier  0.102400\n",
      "Image: book_images/kitty.jpg\n",
      "       label      prob\n",
      "0      tabby  0.173628\n",
      "1  tiger_cat  0.110415\n",
      "2    doormat  0.093447\n",
      "Image: book_images/weaver.jpg\n",
      "       label      prob\n",
      "0    jacamar  0.287284\n",
      "1  bee_eater  0.046768\n",
      "2     bulbul  0.037507\n"
     ]
    }
   ],
   "source": [
    "for i, imgfile in enumerate(imgfiles):\n",
    "    img_df = class_labels.copy()\n",
    "    img_df['prob'] = img_probs[i]\n",
    "    img_df = img_df.sort_values(by='prob', ascending=False)[:3]\n",
    "    print(f'Image: {imgfile}')\n",
    "    print(img_df.reset_index().drop(columns=['idx']))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f9d3843",
   "metadata": {},
   "source": [
    "We see that the model\n",
    "is quite confident about `Flamingo.jpg`, but a little less so for the\n",
    "other images.\n",
    "\n",
    "We end this section with our usual cleanup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "69456669",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "del(cifar_test,\n",
    "    cifar_train,\n",
    "    cifar_dm,\n",
    "    cifar_module,\n",
    "    cifar_logger,\n",
    "    cifar_optimizer,\n",
    "    cifar_trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b33c4d",
   "metadata": {},
   "source": [
    "## IMDB Document Classification\n",
    "We now implement models for sentiment classification (Section 10.4)  on the `IMDB`\n",
    "dataset. As mentioned above code block~8, we are using\n",
    "a preprocessed version of the `IMDB` dataset found in the\n",
    "`keras` package. As `keras` uses `tensorflow`, a different\n",
    "tensor and deep learning library, we have\n",
    "converted the data to be suitable for `torch`. The\n",
    "code used to convert from `keras` is\n",
    "available in the module `ISLP.torch._make_imdb`. It\n",
    "requires some of the `keras` packages to run. These data use a dictionary of size 10,000.\n",
    "\n",
    "We have stored three different representations of the review data for this lab:\n",
    "\n",
    "* `load_tensor()`, a sparse tensor version usable by `torch`;\n",
    "* `load_sparse()`, a sparse matrix version usable by `sklearn`, since we will compare with a lasso fit;\n",
    "* `load_sequential()`, a padded\n",
    "version of the original sequence representation, limited to the last\n",
    "500 words of each review.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "70b3ece2",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([   1,   14,   22,   16,   43,  530,  973, 1622, 1385,   65,  458,\n",
       "       4468], dtype=int32)"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(imdb_seq_train,\n",
    " imdb_seq_test) = load_sequential(root='data/IMDB')\n",
    "padded_sample = np.asarray(imdb_seq_train.tensors[0][0])\n",
    "sample_review = padded_sample[padded_sample > 0][:12]\n",
    "sample_review[:12]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6b73c18",
   "metadata": {},
   "source": [
    "The datasets `imdb_seq_train` and `imdb_seq_test` are\n",
    "both instances of the class `TensorDataset`. The\n",
    "tensors used to construct them can be found in the `tensors` attribute, with\n",
    "the first tensor  the features `X` and the second  the outcome `Y`.\n",
    "We have taken the first row of features and stored it as `padded_sample`. In the preprocessing\n",
    "used to form these data, sequences were padded with 0s in the beginning if they were\n",
    "not long enough, hence we remove this padding by restricting to entries where\n",
    "`padded_sample > 0`. We then provide the first 12 words of the sample review.\n",
    "\n",
    "We can find these words in the `lookup` dictionary from the `ISLP.torch.imdb` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "349b2af7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"<START> this film was just brilliant casting location scenery story direction everyone's\""
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lookup = load_lookup(root='data/IMDB')\n",
    "' '.join(lookup[i] for i in sample_review)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99d5a108",
   "metadata": {},
   "source": [
    "For our first model, we have created a binary feature for each\n",
    "of the 10,000 possible words in the dataset, with an entry of one\n",
    "in the $i,j$ entry if word $j$ appears in review $i$. As most reviews\n",
    "are quite short, such a feature matrix has over 98% zeros. These data\n",
    "are accessed using `load_tensor()` from the `ISLP` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "ffa8da37",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "max_num_workers=10\n",
    "(imdb_train,\n",
    " imdb_test) = load_tensor(root='data/IMDB')\n",
    "imdb_dm = SimpleDataModule(imdb_train,\n",
    "                           imdb_test,\n",
    "                           validation=2000,\n",
    "                           num_workers=min(6, max_num_workers),\n",
    "                           batch_size=512)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e53eea99",
   "metadata": {},
   "source": [
    "We’ll use a two-layer model for our first model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "11eeeabb",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "class IMDBModel(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size):\n",
    "        super(IMDBModel, self).__init__()\n",
    "        self.dense1 = nn.Linear(input_size, 16)\n",
    "        self.activation = nn.ReLU()\n",
    "        self.dense2 = nn.Linear(16, 16)\n",
    "        self.output = nn.Linear(16, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        val = x\n",
    "        for _map in [self.dense1,\n",
    "                     self.activation,\n",
    "                     self.dense2,\n",
    "                     self.activation,\n",
    "                     self.output]:\n",
    "            val = _map(val)\n",
    "        return torch.flatten(val)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943bf8d4",
   "metadata": {},
   "source": [
    "We now instantiate our model and look at a summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "5fe55a29",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "IMDBModel                                [25000, 10003]            [25000]                   --\n",
       "├─Linear: 1-1                            [25000, 10003]            [25000, 16]               160,064\n",
       "├─ReLU: 1-2                              [25000, 16]               [25000, 16]               --\n",
       "├─Linear: 1-3                            [25000, 16]               [25000, 16]               272\n",
       "├─ReLU: 1-4                              [25000, 16]               [25000, 16]               --\n",
       "├─Linear: 1-5                            [25000, 16]               [25000, 1]                17\n",
       "===================================================================================================================\n",
       "Total params: 160,353\n",
       "Trainable params: 160,353\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 4.01\n",
       "===================================================================================================================\n",
       "Input size (MB): 1000.30\n",
       "Forward/backward pass size (MB): 6.60\n",
       "Params size (MB): 0.64\n",
       "Estimated Total Size (MB): 1007.54\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imdb_model = IMDBModel(imdb_test.tensors[0].size()[1])\n",
    "summary(imdb_model,\n",
    "        input_size=imdb_test.tensors[0].size(),\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a49a0145",
   "metadata": {},
   "source": [
    "We’ll again use\n",
    "a smaller learning rate for these data,\n",
    "hence we pass an `optimizer` to the\n",
    "`SimpleModule`. \n",
    "Since the reviews are classified into\n",
    "positive or negative sentiment, we use\n",
    "`SimpleModule.binary_classification()`. {Our use of\n",
    "  `binary_classification()` instead of  `classification()` is\n",
    "  due to some subtlety in how `torchmetrics.Accuracy()` works,\n",
    "as well as the data type of the targets.}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "e1ebbc70",
   "metadata": {},
   "outputs": [],
   "source": [
    "imdb_optimizer = RMSprop(imdb_model.parameters(), lr=0.001)\n",
    "imdb_module = SimpleModule.binary_classification(\n",
    "                         imdb_model,\n",
    "                         optimizer=imdb_optimizer)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18076d4c",
   "metadata": {},
   "source": [
    "Having loaded the datasets into a data module\n",
    "and created a `SimpleModule`, the remaining steps\n",
    "are familiar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "9feddeb2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type              | Params\n",
      "--------------------------------------------\n",
      "0 | model | IMDBModel         | 160 K \n",
      "1 | loss  | BCEWithLogitsLoss | 0     \n",
      "--------------------------------------------\n",
      "160 K     Trainable params\n",
      "0         Non-trainable params\n",
      "160 K     Total params\n",
      "0.641     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1892: PossibleUserWarning: The number of training batches (45) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44936d415c7e4ad6b38cd97c733a90be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
     ]
    }
   ],
   "source": [
    "imdb_logger = CSVLogger('logs', name='IMDB')\n",
    "imdb_trainer = Trainer(deterministic=True,\n",
    "                       max_epochs=30,\n",
    "                       logger=imdb_logger,\n",
    "                       callbacks=[ErrorTracker()])\n",
    "imdb_trainer.fit(imdb_module,\n",
    "                 datamodule=imdb_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28e6658",
   "metadata": {},
   "source": [
    "Evaluating the test error yields roughly 86% accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "887d7dad",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03f177eef9fe4ed090eef15f0a005581",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "      test_accuracy         0.8543599843978882\n",
      "        test_loss           1.1992340087890625\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 1.1992340087890625, 'test_accuracy': 0.8543599843978882}]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_results = imdb_trainer.test(imdb_module, datamodule=imdb_dm)\n",
    "test_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f87545",
   "metadata": {},
   "source": [
    "### Comparison to Lasso\n",
    "We now fit a lasso logistic regression model\n",
    " using `LogisticRegression()` from `sklearn`. Since `sklearn` does not recognize\n",
    "the sparse tensors of `torch`, we use a sparse\n",
    "matrix that is recognized by `sklearn.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "055a9aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "((X_train, Y_train),\n",
    " (X_valid, Y_valid),\n",
    " (X_test, Y_test)) = load_sparse(validation=2000,\n",
    "                                 random_state=0,\n",
    "                                 root='data/IMDB')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8324d03",
   "metadata": {},
   "source": [
    "Similar to what we did in\n",
    "Section 10.9.1,\n",
    "we construct a series of 50 values for the lasso reguralization parameter $\\lambda$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "e3e1e181",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "lam_max = np.abs(X_train.T * (Y_train - Y_train.mean())).max()\n",
    "lam_val = lam_max * np.exp(np.linspace(np.log(1),\n",
    "                                       np.log(1e-4), 50))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73e7d53e",
   "metadata": {},
   "source": [
    "With `LogisticRegression()` the regularization parameter\n",
    "$C$ is specified as the inverse of $\\lambda$. There are several\n",
    "solvers for logistic regression; here we use `liblinear` which\n",
    "works well with the sparse input format. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "dd64350c",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "logit = LogisticRegression(penalty='l1', \n",
    "                           C=1/lam_max,\n",
    "                           solver='liblinear',\n",
    "                           warm_start=True,\n",
    "                           fit_intercept=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e9739a",
   "metadata": {},
   "source": [
    "The path of 50 values takes approximately 40 seconds to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "0bf8b868",
   "metadata": {},
   "outputs": [],
   "source": [
    "coefs = []\n",
    "intercepts = []\n",
    "\n",
    "for l in lam_val:\n",
    "    logit.C = 1/l\n",
    "    logit.fit(X_train, Y_train)\n",
    "    coefs.append(logit.coef_.copy())\n",
    "    intercepts.append(logit.intercept_)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f4fbff",
   "metadata": {},
   "source": [
    "The coefficient and intercepts have an extraneous dimension which can be removed\n",
    "by the `np.squeeze()`  function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "28256828",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "coefs = np.squeeze(coefs)\n",
    "intercepts = np.squeeze(intercepts)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94e746f6",
   "metadata": {},
   "source": [
    "We’ll now make a plot to compare our neural network results with the\n",
    "lasso."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "a5945618",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "fig, axes = subplots(1, 2, figsize=(16, 8), sharey=True)\n",
    "for ((X_, Y_),\n",
    "     data_,\n",
    "     color) in zip([(X_train, Y_train),\n",
    "                    (X_valid, Y_valid),\n",
    "                    (X_test, Y_test)],\n",
    "                    ['Training', 'Validation', 'Test'],\n",
    "                    ['black', 'red', 'blue']):\n",
    "    linpred_ = X_ * coefs.T + intercepts[None,:]\n",
    "    label_ = np.array(linpred_ > 0)\n",
    "    accuracy_ = np.array([np.mean(Y_ == l) for l in label_.T])\n",
    "    axes[0].plot(-np.log(lam_val / X_train.shape[0]),\n",
    "                 accuracy_,\n",
    "                 '.--',\n",
    "                 color=color,\n",
    "                 markersize=13,\n",
    "                 linewidth=2,\n",
    "                 label=data_)\n",
    "axes[0].legend()\n",
    "axes[0].set_xlabel(r'$-\\log(\\lambda)$', fontsize=20)\n",
    "axes[0].set_ylabel('Accuracy', fontsize=20)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dc7e494",
   "metadata": {},
   "source": [
    "Notice the use of `%%capture`, which suppresses the displaying of the partially completed figure. This is useful\n",
    "when making a complex figure, since the steps can be spread across two or more cells.\n",
    "We now add a plot of the lasso accuracy, and display the composed figure by simply entering its name at the end of the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "74704884",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 72844 characters]",
      "text/plain": [
       "<Figure size 1600x800 with 2 Axes>"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imdb_results = pd.read_csv(imdb_logger.experiment.metrics_file_path)\n",
    "summary_plot(imdb_results,\n",
    "             axes[1],\n",
    "             col='accuracy',\n",
    "             ylabel='Accuracy')\n",
    "axes[1].set_xticks(np.linspace(0, 30, 7).astype(int))\n",
    "axes[1].set_ylabel('Accuracy', fontsize=20)\n",
    "axes[1].set_xlabel('Epoch', fontsize=20)\n",
    "axes[1].set_ylim([0.5, 1]);\n",
    "axes[1].axhline(test_results[0]['test_accuracy'],\n",
    "                color='blue',\n",
    "                linestyle='--',\n",
    "                linewidth=3)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7898334d",
   "metadata": {},
   "source": [
    "From the graphs we see that the accuracy of the lasso logistic regression peaks at about $0.88$,  as it does for  the neural network.\n",
    "\n",
    "Once again, we end with a cleanup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "b1ecfca8",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "del(imdb_model,\n",
    "    imdb_trainer,\n",
    "    imdb_logger,\n",
    "    imdb_dm,\n",
    "    imdb_train,\n",
    "    imdb_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0169cd1b",
   "metadata": {},
   "source": [
    "## Recurrent Neural Networks\n",
    "In this lab we fit the models illustrated in\n",
    "Section 10.5."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e7f96cc",
   "metadata": {},
   "source": [
    "### Sequential Models for Document Classification\n",
    "Here we  fit a simple  LSTM RNN for sentiment prediction to\n",
    "the `IMDb` movie-review data, as discussed in Section 10.5.1.\n",
    "For an RNN we use the sequence of words in a document, taking their\n",
    "order into account. We loaded the preprocessed \n",
    "data at the beginning of\n",
    "Section 10.9.5.\n",
    "A script that details the preprocessing can be found in the\n",
    "`ISLP` library. Notably, since more than 90% of the documents\n",
    "had fewer than 500 words, we set the document length to 500. For\n",
    "longer documents, we used the last 500 words, and for shorter\n",
    "documents, we padded the front with blanks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "5503cb3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "imdb_seq_dm = SimpleDataModule(imdb_seq_train,\n",
    "                               imdb_seq_test,\n",
    "                               validation=2000,\n",
    "                               batch_size=300,\n",
    "                               num_workers=min(6, max_num_workers)\n",
    "                               )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "356b76f2",
   "metadata": {},
   "source": [
    "The first layer of the RNN is an embedding layer of size 32, which will be\n",
    "learned during  training. This layer one-hot encodes  each document\n",
    "as a matrix of dimension $500 \\times 10,003$, and then maps these\n",
    "$10,003$ dimensions down to $32$. {The extra 3 dimensions\n",
    "correspond to commonly occurring non-word entries in the reviews.}\n",
    " Since each word is represented by an\n",
    "integer, this is effectively achieved by the creation of an embedding\n",
    "matrix of size $10,003\\times 32$; each of the 500 integers in the\n",
    "document are then mapped to the appropriate 32 real numbers by\n",
    "indexing the appropriate rows of this matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14ddd9a9",
   "metadata": {},
   "source": [
    "The second  layer is an LSTM with 32 units, and the output\n",
    "layer is a single logit for the binary classification task.\n",
    "In the last line of the `forward()` method below,\n",
    "we take the last 32-dimensional output of the LSTM and map it to our response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "8b683641",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "class LSTMModel(nn.Module):\n",
    "    def __init__(self, input_size):\n",
    "        super(LSTMModel, self).__init__()\n",
    "        self.embedding = nn.Embedding(input_size, 32)\n",
    "        self.lstm = nn.LSTM(input_size=32,\n",
    "                            hidden_size=32,\n",
    "                            batch_first=True)\n",
    "        self.dense = nn.Linear(32, 1)\n",
    "    def forward(self, x):\n",
    "        val, (h_n, c_n) = self.lstm(self.embedding(x))\n",
    "        return torch.flatten(self.dense(val[:,-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7903383",
   "metadata": {},
   "source": [
    "We instantiate and take a look at the summary of the model, using the\n",
    "first 10 documents in the corpus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "702aa8de",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "LSTMModel                                [10, 500]                 [10]                      --\n",
       "├─Embedding: 1-1                         [10, 500]                 [10, 500, 32]             320,096\n",
       "├─LSTM: 1-2                              [10, 500, 32]             [10, 500, 32]             8,448\n",
       "├─Linear: 1-3                            [10, 32]                  [10, 1]                   33\n",
       "===================================================================================================================\n",
       "Total params: 328,577\n",
       "Trainable params: 328,577\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 45.44\n",
       "===================================================================================================================\n",
       "Input size (MB): 50.00\n",
       "Forward/backward pass size (MB): 2.56\n",
       "Params size (MB): 1.31\n",
       "Estimated Total Size (MB): 53.87\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lstm_model = LSTMModel(X_test.shape[-1])\n",
    "summary(lstm_model,\n",
    "        input_data=imdb_seq_train.tensors[0][:10],\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f663b95c",
   "metadata": {},
   "source": [
    "The 10,003 is suppressed in the summary, but we see it in the\n",
    "parameter count, since $10,003\\times 32=320,096$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "e48aa272",
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm_module = SimpleModule.binary_classification(lstm_model)\n",
    "lstm_logger = CSVLogger('logs', name='IMDB_LSTM')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "143be7e8",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type              | Params\n",
      "--------------------------------------------\n",
      "0 | model | LSTMModel         | 328 K \n",
      "1 | loss  | BCEWithLogitsLoss | 0     \n",
      "--------------------------------------------\n",
      "328 K     Trainable params\n",
      "0         Non-trainable params\n",
      "328 K     Total params\n",
      "1.314     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1ee0aabd8caf403c93b8c4ba0a1c2507",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=20` reached.\n"
     ]
    }
   ],
   "source": [
    "lstm_trainer = Trainer(deterministic=True,\n",
    "                       max_epochs=20,\n",
    "                       logger=lstm_logger,\n",
    "                       callbacks=[ErrorTracker()])\n",
    "lstm_trainer.fit(lstm_module,\n",
    "                 datamodule=imdb_seq_dm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55dc2b66",
   "metadata": {},
   "source": [
    "The rest is now similar to other networks we have fit. We\n",
    "track the test performance as the network is fit, and see that it attains 85% accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "9272e2ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92a1cacbb1944d74b3f3c0670eb84b63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "      test_accuracy         0.8452399969100952\n",
      "        test_loss           0.7559056878089905\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.7559056878089905, 'test_accuracy': 0.8452399969100952}]"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lstm_trainer.test(lstm_module, datamodule=imdb_seq_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eee55195",
   "metadata": {},
   "source": [
    "We once again show the learning progress, followed by cleanup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "5d9d387e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.5, 1.0)"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 23140 characters]",
      "text/plain": [
       "<Figure size 600x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "lstm_results = pd.read_csv(lstm_logger.experiment.metrics_file_path)\n",
    "fig, ax = subplots(1, 1, figsize=(6, 6))\n",
    "summary_plot(lstm_results,\n",
    "             ax,\n",
    "             col='accuracy',\n",
    "             ylabel='Accuracy')\n",
    "ax.set_xticks(np.linspace(0, 20, 5).astype(int))\n",
    "ax.set_ylabel('Accuracy')\n",
    "ax.set_ylim([0.5, 1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "7c6d1072",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "del(lstm_model,\n",
    "    lstm_trainer,\n",
    "    lstm_logger,\n",
    "    imdb_seq_dm,\n",
    "    imdb_seq_train,\n",
    "    imdb_seq_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d179385",
   "metadata": {},
   "source": [
    "### Time Series Prediction\n",
    "We now show how to fit the models in Section 10.5.2\n",
    "for  time series prediction.\n",
    "We first load and standardize the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "8e2d1d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "NYSE = load_data('NYSE')\n",
    "cols = ['DJ_return', 'log_volume', 'log_volatility']\n",
    "X = pd.DataFrame(StandardScaler(\n",
    "                     with_mean=True,\n",
    "                     with_std=True).fit_transform(NYSE[cols]),\n",
    "                 columns=NYSE[cols].columns,\n",
    "                 index=NYSE.index)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad4a748",
   "metadata": {},
   "source": [
    "Next we set up the lagged versions of the data, dropping\n",
    "any rows with missing values using the `dropna()`  method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "6f6d0e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "for lag in range(1, 6):\n",
    "    for col in cols:\n",
    "        newcol = np.zeros(X.shape[0]) * np.nan\n",
    "        newcol[lag:] = X[col].values[:-lag]\n",
    "        X.insert(len(X.columns), \"{0}_{1}\".format(col, lag), newcol)\n",
    "X.insert(len(X.columns), 'train', NYSE['train'])\n",
    "X = X.dropna()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a924d2",
   "metadata": {},
   "source": [
    "Finally, we extract the response, training indicator, and drop the current day’s `DJ_return` and\n",
    "`log_volatility` to predict only from previous day’s data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "d3c961ec",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['DJ_return_1', 'log_volume_1', 'log_volatility_1', 'DJ_return_2',\n",
       "       'log_volume_2', 'log_volatility_2', 'DJ_return_3', 'log_volume_3',\n",
       "       'log_volatility_3', 'DJ_return_4', 'log_volume_4', 'log_volatility_4',\n",
       "       'DJ_return_5', 'log_volume_5', 'log_volatility_5'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y, train = X['log_volume'], X['train']\n",
    "X = X.drop(columns=['train'] + cols)\n",
    "X.columns\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c5f8d2f",
   "metadata": {},
   "source": [
    "We first fit a simple linear model and compute the $R^2$ on the test data using\n",
    "the `score()`  method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "cb89da88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.4128912938562521"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M = LinearRegression()\n",
    "M.fit(X[train], Y[train])\n",
    "M.score(X[~train], Y[~train])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b92425ab",
   "metadata": {},
   "source": [
    "We refit this model, including the factor variable `day_of_week`.\n",
    "For a categorical series in `pandas`, we can form the indicators\n",
    "using the `get_dummies()`  method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "55d921d7",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "X_day = pd.merge(X, \n",
    "                 pd.get_dummies(NYSE['day_of_week']),\n",
    "                 on='date')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afab736a",
   "metadata": {},
   "source": [
    " Note that we do not have\n",
    "to reinstantiate the linear regression model\n",
    "as its `fit()` method accepts a design matrix and a response directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "4e87eeb9",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.45633025715446285"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M.fit(X_day[train], Y[train])\n",
    "M.score(X_day[~train], Y[~train])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43309e3b",
   "metadata": {},
   "source": [
    "This model achieves an $R^2$ of about 46%."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93c61d80",
   "metadata": {},
   "source": [
    "To fit the RNN, we must reshape the data, as it will expect 5 lagged\n",
    "versions of each feature as indicated by the  `input_shape` argument\n",
    "to the layer `nn.RNN()`  below. We first\n",
    "ensure the columns of  our data frame are such that a reshaped\n",
    "matrix will have the variables correctly lagged. We use the\n",
    "`reindex()`  method to do this.\n",
    "\n",
    "For an input shape `(5,3)`, each row represents a lagged version of the three variables.\n",
    "The `nn.RNN()` layer also expects the first row of each\n",
    "observation to be earliest in time, so we must reverse the current order.  Hence we loop over\n",
    "`range(5,0,-1)` below, which is\n",
    "an example of using a  `slice()`  to index\n",
    "iterable objects. The general notation is `start:end:step`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "3689b318",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['DJ_return_5', 'log_volume_5', 'log_volatility_5', 'DJ_return_4',\n",
       "       'log_volume_4', 'log_volatility_4', 'DJ_return_3', 'log_volume_3',\n",
       "       'log_volatility_3', 'DJ_return_2', 'log_volume_2', 'log_volatility_2',\n",
       "       'DJ_return_1', 'log_volume_1', 'log_volatility_1'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ordered_cols = []\n",
    "for lag in range(5,0,-1):\n",
    "    for col in cols:\n",
    "        ordered_cols.append('{0}_{1}'.format(col, lag))\n",
    "X = X.reindex(columns=ordered_cols)\n",
    "X.columns\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dbcf9f7",
   "metadata": {},
   "source": [
    "We now reshape the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "5633e521",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6046, 5, 3)"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_rnn = X.to_numpy().reshape((-1,5,3))\n",
    "X_rnn.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd3ae3ff",
   "metadata": {},
   "source": [
    "By specifying the first size as -1, `numpy.reshape()` deduces its size based on the remaining arguments.\n",
    "\n",
    "Now we are ready to proceed with the RNN, which uses 12 hidden units, and  10%\n",
    "dropout.\n",
    "After passing through the RNN,  we extract the final time point as `val[:,-1]`\n",
    "in `forward()` below. This gets passed through a 10% dropout and then flattened through\n",
    "a linear layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "979a6066",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NYSEModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NYSEModel, self).__init__()\n",
    "        self.rnn = nn.RNN(3,\n",
    "                          12,\n",
    "                          batch_first=True)\n",
    "        self.dense = nn.Linear(12, 1)\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "    def forward(self, x):\n",
    "        val, h_n = self.rnn(x)\n",
    "        val = self.dense(self.dropout(val[:,-1]))\n",
    "        return torch.flatten(val)\n",
    "nyse_model = NYSEModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14671ab2",
   "metadata": {},
   "source": [
    "We  fit the model in a similar fashion to previous networks. We\n",
    "supply the `fit` function with test data as validation data, so that when\n",
    "we monitor its progress and plot the history function we can see the\n",
    "progress on the test data. Of course we should not use this as a basis for\n",
    "early stopping, since then the test performance would be biased.\n",
    "\n",
    "We form the training dataset similar to\n",
    "our `Hitters` example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "3528f4e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = []\n",
    "for mask in [train, ~train]:\n",
    "    X_rnn_t = torch.tensor(X_rnn[mask].astype(np.float32))\n",
    "    Y_t = torch.tensor(Y[mask].astype(np.float32))\n",
    "    datasets.append(TensorDataset(X_rnn_t, Y_t))\n",
    "nyse_train, nyse_test = datasets\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a731099",
   "metadata": {},
   "source": [
    "Following our usual pattern, we inspect the summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "795c283e",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torchinfo/torchinfo.py:477: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  action_fn=lambda data: sys.getsizeof(data.storage()),\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/torch/storage.py:665: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return super().__sizeof__() + self.nbytes()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Input Shape               Output Shape              Param #\n",
       "===================================================================================================================\n",
       "NYSEModel                                [1770, 5, 3]              [1770]                    --\n",
       "├─RNN: 1-1                               [1770, 5, 3]              [1770, 5, 12]             204\n",
       "├─Dropout: 1-2                           [1770, 12]                [1770, 12]                --\n",
       "├─Linear: 1-3                            [1770, 12]                [1770, 1]                 13\n",
       "===================================================================================================================\n",
       "Total params: 217\n",
       "Trainable params: 217\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (M): 1.83\n",
       "===================================================================================================================\n",
       "Input size (MB): 0.11\n",
       "Forward/backward pass size (MB): 0.86\n",
       "Params size (MB): 0.00\n",
       "Estimated Total Size (MB): 0.97\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(nyse_model,\n",
    "        input_data=X_rnn_t,\n",
    "        col_names=['input_size',\n",
    "                   'output_size',\n",
    "                   'num_params'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "650e0ec7",
   "metadata": {},
   "source": [
    "\\newpage\n",
    "We again put the two datasets into a data module, with a\n",
    "batch size of 64."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "31640121",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "nyse_dm = SimpleDataModule(nyse_train,\n",
    "                           nyse_test,\n",
    "                           num_workers=min(4, max_num_workers),\n",
    "                           validation=nyse_test,\n",
    "                           batch_size=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1de7ae89",
   "metadata": {},
   "source": [
    "We run some data through our model to be sure the sizes match up correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "af6795f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64]) torch.Size([64])\n",
      "torch.Size([64]) torch.Size([64])\n",
      "torch.Size([64]) torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "for idx, (x, y) in enumerate(nyse_dm.train_dataloader()):\n",
    "    out = nyse_model(x)\n",
    "    print(y.size(), out.size())\n",
    "    if idx >= 2:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eecb6d9c",
   "metadata": {},
   "source": [
    "We follow our previous example for setting up a trainer for a\n",
    "regression problem, requesting the $R^2$ metric\n",
    "to be be computed at each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "0d04344c",
   "metadata": {},
   "outputs": [],
   "source": [
    "nyse_optimizer = RMSprop(nyse_model.parameters(),\n",
    "                         lr=0.001)\n",
    "nyse_module = SimpleModule.regression(nyse_model,\n",
    "                                      optimizer=nyse_optimizer,\n",
    "                                      metrics={'r2':R2Score()})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a10444d",
   "metadata": {},
   "source": [
    "Fitting the model should by now be familiar.\n",
    "The results on the test data are very similar to the linear AR model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "485d7bb1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/jonathantaylor/anaconda3/envs/islp_freeze_39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name  | Type      | Params\n",
      "------------------------------------\n",
      "0 | model | NYSEModel | 217   \n",
      "1 | loss  | MSELoss   | 0     \n",
      "------------------------------------\n",
      "217       Trainable params\n",
      "0         Non-trainable params\n",
      "217       Total params\n",
      "0.001     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf9eb739cc954760813987d94e79fb8e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce04cc5cab594022a6b0e9a1f429ba10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a45ddcb60d7b4bd69dfaa5e0338d7dcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c95a6371f5af4aeaa0a58a5f5b85db44",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89e3e03d38994cf1a1093722c2c4fe1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "951e5cbfbb804e779bbbaf4b4a351ae2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d3d473df90f4dd3aeb829d52d35d641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7395df14478d4bed859a3f9cf8ccbead",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef61c7b225144a7398cff81780b0e60b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "404284c1c9ea41628cb9bb6f4d031cd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b651d8b168ee452a98c397148c7d1aea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c66168b0e85424c8abe2446f7f9e5a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6998335a07b341f39fd4ebf1c93d162f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "076763c8702042028a8d216e15fef774",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "abb0d6faecc24e059b3e68671dc38ad5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "140dcdd80e2443619f46d24a17ed9d50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "991f4ae3b2b249848b57f5015f5fe6c7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95da57cb36fb4ebaa0d52fa94d6c4b32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d6915797540439da8b0448c860eb604",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5855d2022474bb7923ec7915138699b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "209f7ce81cd24ef2976c054769619f84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1924729f28a84ff3b7c9c8607cea1de7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "58b7467b4a7c4b10aee2d9bf3eb3fafc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce9b0241a7384110a1b0291f01f3780b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ec8d539b2ab4f58beb3b5f49764fff3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41c8e961dac34415a4ccceb8babbcf01",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ac4d97ce9634218bdf232cca63d6069",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ffea75eb7ccc4af3b66af673c485401b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80c06a18bd4e4b159ff4f0bd1b333994",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "911f58a64ce24c9b8825fdd6372580eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37e309ed15c34ebf9ad329024eda3c23",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59a073d533f0447fa920155ad6e9aa1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5175feb3fd374cd89de5c62e3d3b1b08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce0ca0dae31c498ea244ccdcfae501e7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd510768c59f427787f55d787ea8e5cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95dfb50b714949bf8e28d8bbe00736bf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "437d2e7134b24877afc41e5355fb686c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fc6d6a6db814fcca6cdc78e63f3f816",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "45a341642e2d4080ad9554826fe775f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b5a45975a73c4d3f9b6d19dc27d249f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "742d6b44d2a14afba4b98a4e4ce75e2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60e2119e8ae348ab9368218e4fef8c7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aac8c13ba5864e6cb2d62a2c6a5dde86",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0022bae4698c4f52b7515abb1a444d4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ff5051151de94de98de6be6b6ee5c6b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38bf027d00494f90a40bb178fe6db007",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6cf38392f8c448883140a325ca52878",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae0806202fe441d68d48773578f252c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a788a24e8f843e18e6013d7d76d1041",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b46aad0207543afad9d9ac4670605c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e90e2cf3add47c7958f4025508c0f40",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3750b847af554f3ebe96677e2d2580d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd50a1b812d04368968cf4c4ecf3f639",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5c7d28ab7c44c0aae170b05a3c6e923",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c32ddfe9a7f44ee8a1485716dd209ca1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47a99671fd634e1799addd33bfaca5ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98482279444c4db9bc48fbfe02dfb7dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a56dbc5229434698a43139dc4ff3a460",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b52732b1d1e94f46a1cd037b29d27ce0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f1264a1828b4d55b74f0dccd97919c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59bf88e571724091ab98f7258dbe4780",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "632942952df344ee98d2423705e060c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "152694e1fe214fe899cbeba080d5fec7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92e87ecd035e4a98af8940873c217a08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5cfb6d3324224105bb51a11ee3610e99",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b11d5fcedba4da68bec988d9c1af1a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39bf95953cd84ff1b5206301be0caba9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4507f5f3b42c4f79815e72a55cd57101",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c47d400fe5ac44098a50f94c59223d71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b580013d7f0447a96878fa0982eda94",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "49e3e8b33931437bb03110aacac4c11d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82e34b2ab25a4d52bd58dd791fab350b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0233c4562c014826845209456bb8d5f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "        test_loss           0.6207807064056396\n",
      "         test_r2            0.41084879636764526\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.6207807064056396, 'test_r2': 0.41084879636764526}]"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nyse_trainer = Trainer(deterministic=True,\n",
    "                       max_epochs=200,\n",
    "                       callbacks=[ErrorTracker()])\n",
    "nyse_trainer.fit(nyse_module,\n",
    "                 datamodule=nyse_dm)\n",
    "nyse_trainer.test(nyse_module,\n",
    "                  datamodule=nyse_dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbd13ee3",
   "metadata": {},
   "source": [
    "We could also fit a model without the `nn.RNN()` layer by just\n",
    "using a `nn.Flatten()` layer instead. This would be a nonlinear AR model. If in addition we excluded the \n",
    "hidden layer, this would be equivalent to our earlier linear AR model.  \n",
    "\n",
    "Instead we will fit a nonlinear AR model using the feature set `X_day` that includes the `day_of_week` indicators.\n",
    "To do so, we\n",
    "must first create our test and training datasets and a corresponding\n",
    "data module. This may seem a little burdensome, but is part of the\n",
    "general pipeline for `torch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "91d6633a",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = []\n",
    "for mask in [train, ~train]:\n",
    "    X_day_t = torch.tensor(\n",
    "                   np.asarray(X_day[mask]).astype(np.float32))\n",
    "    Y_t = torch.tensor(np.asarray(Y[mask]).astype(np.float32))\n",
    "    datasets.append(TensorDataset(X_day_t, Y_t))\n",
    "day_train, day_test = datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5cf941c",
   "metadata": {},
   "source": [
    "Creating a data module follows a familiar pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "87137503",
   "metadata": {},
   "outputs": [],
   "source": [
    "day_dm = SimpleDataModule(day_train,\n",
    "                          day_test,\n",
    "                          num_workers=min(4, max_num_workers),\n",
    "                          validation=day_test,\n",
    "                          batch_size=64)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2575f374",
   "metadata": {},
   "source": [
    "We build a `NonLinearARModel()` that takes as input the 20 features and a hidden layer  with 32 units. The remaining steps are familiar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "3cf09a55",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NonLinearARModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NonLinearARModel, self).__init__()\n",
    "        self._forward = nn.Sequential(nn.Flatten(),\n",
    "                                      nn.Linear(20, 32),\n",
    "                                      nn.ReLU(),\n",
    "                                      nn.Dropout(0.5),\n",
    "                                      nn.Linear(32, 1))\n",
    "    def forward(self, x):\n",
    "        return torch.flatten(self._forward(x))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "930576f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "nl_model = NonLinearARModel()\n",
    "nl_optimizer = RMSprop(nl_model.parameters(),\n",
    "                           lr=0.001)\n",
    "nl_module = SimpleModule.regression(nl_model,\n",
    "                                        optimizer=nl_optimizer,\n",
    "                                        metrics={'r2':R2Score()})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e3043c",
   "metadata": {},
   "source": [
    "We continue with the usual training steps, fit the model,\n",
    "and evaluate the test error. We see the test $R^2$ is a slight improvement over the linear AR model that also includes `day_of_week`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "b9ae8d0e",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | NonLinearARModel | 705   \n",
      "1 | loss  | MSELoss          | 0     \n",
      "-------------------------------------------\n",
      "705       Trainable params\n",
      "0         Non-trainable params\n",
      "705       Total params\n",
      "0.003     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e278bfc60e6404f98175752ac8f6482",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "        test_loss           0.5648325681686401\n",
      "         test_r2            0.4639463424682617\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.5648325681686401, 'test_r2': 0.4639463424682617}]"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nl_trainer = Trainer(deterministic=True,\n",
    "                         max_epochs=20,\n",
    "                         callbacks=[ErrorTracker()])\n",
    "nl_trainer.fit(nl_module, datamodule=day_dm)\n",
    "nl_trainer.test(nl_module, datamodule=day_dm) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd356e46",
   "metadata": {},
   "source": [
    "              \n",
    "\n",
    " \n"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
