notebooks/model_tutorial/0_baseline_CNN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Baseline CNN model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this notebook we provide an explanation of the baseline CNN model, including: \n",
    "- a written description of the model (with references for further information)\n",
    "- a code demo exploring different parts of the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Baseline CNN model description"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### The model has two main parts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The baseline CNN model (which is mostly based on [this work](https://openreview.net/forum?id=Tp7kI90Htd)) is constructed from two main parts:\n",
    "- **core**: the core aims to (nonlinearly) extract features that are common between neurons. That is, we assume there exist a set of features that all neurons use but combine them in their own unique way.\n",
    "- **readout**: once the core extracts the feautures, then a neuron reads out from those features by simply linearly combining those features into a single value. Finally, by passing this single value through a final nonlinarity (in this case `ELU() + 1`) we make sure that the model output is positive and we get the inferred firing rate of the neuron."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Learning where neurons \"look\" 👀"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "From experimental evidence, we know that neurons are sensitive to a limited area in the visual field - this is referred to as neuron's Receptive Field (RF). Knowing this, the readout is equipped with a mechanism that allows the model to learn where the neuron is \"looking\" in the visual field. In other words, the model learns the RF position of the neuron, allowing it to pick a specific spatial position from the core's output and then linearly combine the features along the channel dimension. This significantly reduces the number of parameters in the readout."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### RF as a function of the cortex positions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "While the neuron's receptive field can be defined as a model parameter and be learned during training, we can be further inspired by the experimental evidence: neuron's that are located close to each other on the cortex, their RFs are also close in the visual field. To this end, we equip the readout with an additional module called **readout position network** which learns to map cortical positions to RF positions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Accounting for RF shifts due to eye movement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, we account for shifts in the RF positions due to eye movement with yet another module called **shifter**. The shifter module takes pupil position (2d vector) as input and outputs a 2d vector which is used to globally shift all neurons' RF position."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### References\n",
    "- [**Paper**: Lurz, K. K., Bashiri, M., Willeke, K., Jagadish, A. K., Wang, E., Walker, E. Y., ... & Sinz, F. H. (2021). Generalization in data-driven models of primary visual cortex. BioRxiv, 2020-10.](https://www.biorxiv.org/content/10.1101/2020.10.05.326256v2)\n",
    "- [**This Video**](https://youtu.be/xwLMO8nVvxs?t=220) (which is a talk explaining the above paper) also explains the readout."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Baseline CNN model exploration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from nnfabrik.builder import get_data\n",
    "\n",
    "device = \"cuda\"\n",
    "random_seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Instantiate DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To initialize the model we use a *model function* which requires the dataloader as an input argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# loading the SENSORIUM+ dataset\n",
    "filenames = ['../data/static27204-5-13-GrayImageNet-94c6ff995dac583098847cfecd43e7b6.zip', ]\n",
    "\n",
    "dataset_fn = 'sensorium.datasets.static_loaders'\n",
    "dataset_config = {'paths': filenames,\n",
    "                 'normalize': True,\n",
    "                 'include_behavior': False,\n",
    "                 'include_eye_position': True,\n",
    "                 'batch_size': 32,\n",
    "                 'scale':1,\n",
    "                 }\n",
    "\n",
    "dataloaders = get_data(dataset_fn, dataset_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Import the model function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from sensorium.models import stacked_core_full_gauss_readout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's have a quick look at the inputs arguments of the model function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m\n",
       "\u001b[0mstacked_core_full_gauss_readout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdataloaders\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mseed\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mhidden_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m32\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minput_kern\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m13\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mhidden_kern\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mlayers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mgamma_input\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m15.5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mskip\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mfinal_nonlinearity\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mmomentum\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.9\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mpad_input\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mbatch_norm\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mhidden_dilation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mlaplace_padding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minput_regularizer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'LaplaceL2norm'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0muse_avg_reg\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minit_mu_range\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minit_sigma\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1.0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mreadout_bias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mgamma_readout\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0melu_offset\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mstack\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdepth_separable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mlinear\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mgauss_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'full'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mgrid_mean_predictor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mattention_conv\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mshifter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mshifter_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'MLP'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minput_channels_shifter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mhidden_channels_shifter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mshift_layers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mgamma_shifter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mshifter_bias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mhidden_padding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mcore_bias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Model class of a stacked2dCore (from neuralpredictors) and a pointpooled (spatial transformer) readout\n",
       "\n",
       "Args:\n",
       "    dataloaders: a dictionary of dataloaders, one loader per session\n",
       "        in the format {'data_key': dataloader object, .. }\n",
       "    seed: random seed\n",
       "    grid_mean_predictor: if not None, needs to be a dictionary of the form\n",
       "        {\n",
       "        'type': 'cortex',\n",
       "        'input_dimensions': 2,\n",
       "        'hidden_layers':0,\n",
       "        'hidden_features':20,\n",
       "        'final_tanh': False,\n",
       "        }\n",
       "        In that case the datasets need to have the property `neurons.cell_motor_coordinates`\n",
       "    share_features: whether to share features between readouts. This requires that the datasets\n",
       "        have the properties `neurons.multi_match_id` which are used for matching. Every dataset\n",
       "        has to have all these ids and cannot have any more.\n",
       "    all other args: See Documentation of Stacked2dCore in neuralpredictors.layers.cores and\n",
       "        PointPooled2D in neuralpredictors.layers.readouts\n",
       "\n",
       "Returns: An initialized model which consists of model.core and model.readout\n",
       "\u001b[0;31mFile:\u001b[0m      /project/sensorium/models/models.py\n",
       "\u001b[0;31mType:\u001b[0m      function\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "stacked_core_full_gauss_readout?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Specify (some of) the input arguments to initialize the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "grid_mean_predictor = {\n",
    "    'type': 'cortex',\n",
    "    'input_dimensions': 2,\n",
    "    'hidden_layers': 1,\n",
    "    'hidden_features': 30,\n",
    "    'final_tanh': True\n",
    "}\n",
    "\n",
    "model_config = {\n",
    "    # core args\n",
    "    'input_kern': 9,\n",
    "    'hidden_kern': 7,\n",
    "    'hidden_channels': 64,\n",
    "    'layers': 4,\n",
    "    'depth_separable': True,\n",
    "    'stack': -1,\n",
    "    'gamma_input': 6.3831,\n",
    "    # readout args\n",
    "    'gamma_readout': 0.0076,\n",
    "    'grid_mean_predictor': grid_mean_predictor,\n",
    "    'gauss_type': 'full',\n",
    "    'shifter': True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "With the above model config we are defining a model which:\n",
    "- has four convolution laters where each layer\n",
    "    - is a depth-seperable convolution (`depth_separable=True`)\n",
    "    - has a kernel size of 7 (`hidden_kern=7`), with the exception of the first layer which has a kernel size of 9 (`input_kern=9`)\n",
    "    - outputs an activation tensor with 64 channels (`hidden_channels=64`)\n",
    "- only uses the last layer (`stack=-1`) as the final output of the core (other options results in stacking the outputs of multiple layers)\n",
    "- uses the cortex positions (x and y) of the neurons to infer their receptive field position (specified with the `grid_mean_predictor`)\n",
    "- uses the pupil center to shift the neurons' receptive positions globally (they will all be shifted the same way) depending on the pupil position of the subject (`shifter=True`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Instantiate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model = stacked_core_full_gauss_readout(dataloaders, random_seed, **model_config).to(device)\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FiringRateEncoder(\n",
       "  (core): Stacked2dCore(\n",
       "    (_input_weights_regularizer): LaplaceL2norm(\n",
       "      (laplace): Laplace()\n",
       "    )\n",
       "    (features): Sequential(\n",
       "      (layer0): Sequential(\n",
       "        (conv): Conv2d(1, 64, kernel_size=(9, 9), stride=(1, 1), bias=False)\n",
       "        (norm): BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer1): Sequential(\n",
       "        (ds_conv): DepthSeparableConv2d(\n",
       "          (in_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (spatial_conv): Conv2d(64, 64, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=64, bias=False)\n",
       "          (out_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        )\n",
       "        (norm): BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (ds_conv): DepthSeparableConv2d(\n",
       "          (in_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (spatial_conv): Conv2d(64, 64, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=64, bias=False)\n",
       "          (out_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        )\n",
       "        (norm): BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer3): Sequential(\n",
       "        (ds_conv): DepthSeparableConv2d(\n",
       "          (in_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (spatial_conv): Conv2d(64, 64, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=64, bias=False)\n",
       "          (out_depth_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        )\n",
       "        (norm): BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "    )\n",
       "  ) [Stacked2dCore regularizers: gamma_hidden = 0|gamma_input = 6.3831|skip = 0]\n",
       "  \n",
       "  (readout): MultipleFullGaussian2d(\n",
       "    (27204-5-13): full FullGaussian2d (64 x 136 x 248 -> 7538) with bias, with predicted grid  -> Sequential(\n",
       "      (0): Linear(in_features=2, out_features=30, bias=True)\n",
       "      (1): ELU(alpha=1.0)\n",
       "      (2): Linear(in_features=30, out_features=2, bias=True)\n",
       "      (3): Tanh()\n",
       "    )\n",
       "    \n",
       "  )\n",
       "  (shifter): MLPShifter(\n",
       "    (27204-5-13): MLP(\n",
       "      (mlp): Sequential(\n",
       "        (0): Linear(in_features=2, out_features=5, bias=True)\n",
       "        (1): Tanh()\n",
       "        (2): Linear(in_features=5, out_features=5, bias=True)\n",
       "        (3): Tanh()\n",
       "        (4): Linear(in_features=5, out_features=2, bias=True)\n",
       "        (5): Tanh()\n",
       "      )\n",
       "    ) [MLP regularizers: ]\n",
       "    \n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Looking at the model we can see three top modules:\n",
    "- core: `model.core`\n",
    "- readout: `model.readout`\n",
    "- shifter: `model.shifter`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "**Note** that when you display the readout or the shifter modules you get a `ModuleDict` as opposed to the core which is just a `Module`. The model is designed this way to allow the user to share the core between multiple datasets. While the features (i.e. output of the core) can be shared between different neurons and different subjects/sessions, other information are most likely unique to each dataset. Therefore, we keep the readout and shifter dataset specific. Each key (aka `data_key`) in the readout (or shifter) corresponds to a specific dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we are going to go through each model and explore them individually. But before doing that let's get a batch of data:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Get a single batch from the dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data_key = '27204-5-13'\n",
    "batch = next(iter(dataloaders[\"train\"][data_key]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['images', 'responses', 'pupil_center'])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch._asdict().keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 1, 144, 256])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 7538])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.responses.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 2])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.pupil_center.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "core_output = model.core(batch.images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 64, 136, 248])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core_output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We see that the output of the core has:\n",
    "- **64** channels as specified by the `hidden_channels` argument\n",
    "- a height of **136** which is `input_h - input_kern + 1` <-> **144 - 9 + 1 = 136**. Note that this also implies that the rest of the convolutional layers preserve the dimensions (i.e. `padding=same`)\n",
    "- a width of **248** which is `input_w - input_kern + 1` <-> **256 - 9 + 1 = 248**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Readout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "readout_output = model.readout[data_key](core_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 7538])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "readout_output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And here we we have the predicted firing rate of all the neurons (n=7538) for all the images (n=32) in this batch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Where is the shifter then?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The shifter takes the eye position (can be assessed in a variable called `pupil_center`) and outputs a global shift which is then used as an input to the readout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "shifter_output = model.shifter[data_key](batch.pupil_center)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 2])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shifter_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "readout_output_shifted = model.readout[data_key](core_output, shift=shifter_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 7538])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "readout_output_shifted.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can check whether this results in a different output from the readout:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.equal(readout_output, readout_output_shifted)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### References\n",
    "- Code for the [model function](https://github.com/sinzlab/sensorium/blob/8660c0c925b3944e723637db4725083f84ee28c3/sensorium/models/models.py#L17)\n",
    "- Code for the [core Module](https://github.com/sinzlab/neuralpredictors/blob/0d3d793cc0e1f55ec61c5f9f7a98318b5241a2e9/neuralpredictors/layers/cores/conv2d.py#L27)\n",
    "- Code for the [readout Module](https://github.com/sinzlab/neuralpredictors/blob/0d3d793cc0e1f55ec61c5f9f7a98318b5241a2e9/neuralpredictors/layers/readouts/gaussian.py#L210)\n",
    "- Code for the [shifter module](https://github.com/sinzlab/neuralpredictors/blob/0d3d793cc0e1f55ec61c5f9f7a98318b5241a2e9/neuralpredictors/layers/shifters/mlp.py#L13)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}