notebooks/benchmarks_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f713957f-16a1-4f73-9a0f-c340d5eb94da",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "33513917-9ec0-4b02-aba8-99bae0cdf243",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if imports of sensorium do not work \n",
    "# !ln -s ../sensorium\n",
    "\n",
    "import torch\n",
    "from nnfabrik.utility.nn_helpers import set_random_seed\n",
    "set_random_seed(seed)\n",
    "\n",
    "from sensorium.datasets.mouse_video_loaders import mouse_video_loader\n",
    "from sensorium.utility.scores import get_correlations\n",
    "from nnfabrik.builder import get_trainer\n",
    "from sensorium.models.make_model import make_video_model\n",
    "\n",
    "device = 'cuda'\n",
    "torch.cuda.set_device(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e4e87365-018a-4fe8-b54e-d49409ff95a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data..\n",
      "Data loaded\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "    paths to data folders like\n",
    "    paths = [\n",
    "    '/my_path/dynamic29515-10-12-Video-9b4f6a1a067fe51e15306b9628efea20/',\n",
    "    '/my_path/dynamic29623-4-9-Video-9b4f6a1a067fe51e15306b9628efea20/',\n",
    "    '/my_path/dynamic29647-19-8-Video-9b4f6a1a067fe51e15306b9628efea20/',\n",
    "    '/my_path/dynamic29712-5-9-Video-9b4f6a1a067fe51e15306b9628efea20/',\n",
    "    '/my_path/dynamic29755-2-8-Video-9b4f6a1a067fe51e15306b9628efea20/'\n",
    "    ]\n",
    "    \n",
    "    the \"/\" at the end is important\n",
    "'''\n",
    "\n",
    "print(\"Loading data..\")\n",
    "data_loaders = mouse_video_loader(\n",
    "    paths=paths,\n",
    "    batch_size=8,\n",
    "    scale=1,\n",
    "    max_frame=None,\n",
    "    frames=80, # frames has to be > 50.\n",
    "    offset=-1,\n",
    "    include_behavior=True,\n",
    "    include_pupil_centers=True,\n",
    "    cuda=device!='cpu',\n",
    ")\n",
    "print('Data loaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "157ff411-7b6b-44cb-b4c2-c440354849f7",
   "metadata": {},
   "source": [
    "## GRU Benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "742e7d9f-b197-4e17-b537-da280fd2152e",
   "metadata": {},
   "outputs": [],
   "source": [
    "equivar_2D_core_dict = dict(\n",
    "    input_channels=3,\n",
    "    hidden_channels=8,\n",
    "    input_kern=9,\n",
    "    hidden_kern=7,\n",
    "    layers=4,\n",
    "    num_rotations=8,\n",
    "    gamma_input=500,\n",
    "    skip=0,\n",
    "    pad_input=False,\n",
    "    final_nonlinearity=False,\n",
    "    bias=True,\n",
    "    momentum=0.9,\n",
    "    batch_norm=True,\n",
    "    hidden_dilation=1,\n",
    "    laplace_padding=None,\n",
    "    input_regularizer=\"LaplaceL2norm\",\n",
    "    stack=-1,\n",
    "    depth_separable=False,\n",
    "    linear=False,\n",
    "    attention_conv=False,\n",
    "    hidden_padding=None,\n",
    "    use_avg_reg=False,\n",
    "    final_batchnorm_scale=True,\n",
    "    gamma_hidden=500_000,\n",
    ")\n",
    "\n",
    "gru_dict = dict(\n",
    "    # input channels should be the last hidden channels from the core_dict\n",
    "    input_channels=64,\n",
    "    # rec channels should be the input channels to the readouts\n",
    "    rec_channels=64,\n",
    "    input_kern=9,\n",
    "    rec_kern=9,\n",
    "    gamma_rec=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0b538372-2034-45cc-a8b9-5023e9c51e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "shifter_dict = dict(\n",
    "    gamma_shifter=0,\n",
    "    shift_layers=3,\n",
    "    input_channels_shifter=2,\n",
    "    hidden_channels_shifter=5,\n",
    ")\n",
    "\n",
    "\n",
    "readout_dict = dict(\n",
    "    bias=True,\n",
    "    init_mu_range=0.2,\n",
    "    init_sigma=1.0,\n",
    "    gamma_readout=0.0,\n",
    "    gauss_type='full',\n",
    "    grid_mean_predictor={\n",
    "        'type': 'cortex',\n",
    "        'input_dimensions': 2,\n",
    "        'hidden_layers': 1,\n",
    "        'hidden_features': 30,\n",
    "        'final_tanh': True\n",
    "    },\n",
    "    share_features=False,\n",
    "    share_grid=False,\n",
    "    shared_match_ids=None,\n",
    "    gamma_grid_dispersion=0.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "efe2434b-c3e0-416a-8459-cf118cb8e540",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/readouts/base.py:74: UserWarning: Use of 'gamma_readout' is deprecated. Use 'feature_reg_weight' instead. If 'feature_reg_weight' is defined, 'gamma_readout' is ignored\n",
      "  warnings.warn(\n",
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/rnn_modules/gru_module.py:18: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.\n",
      "  xavier_normal(m.weight.data)\n",
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/rnn_modules/gru_module.py:20: UserWarning: nn.init.constant is now deprecated in favor of nn.init.constant_.\n",
      "  init.constant(m.bias.data, 0.0)\n"
     ]
    }
   ],
   "source": [
    "gru_2d_model_equivariant = make_video_model(data_loaders,\n",
    "                 seed,\n",
    "                 core_dict=equivar_2D_core_dict,\n",
    "                 core_type='2D_equivariant',\n",
    "                 readout_dict=readout_dict.copy(),\n",
    "                 readout_type='gaussian',               \n",
    "                 use_gru=True,\n",
    "                 gru_dict=gru_dict,\n",
    "                 use_shifter=True,\n",
    "                 shifter_dict=shifter_dict,\n",
    "                 shifter_type='MLP',\n",
    "                                            \n",
    "                 # todo - put this to True if you are using deeplake\n",
    "                 # first connections to deeplake may take up for 10 mins\n",
    "                 deeplake_ds=False,\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e2c5cd67-08cb-4d07-9f9f-7e1258608136",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VideoFiringRateEncoder(\n",
       "  (core): RotationEquivariant2dCore(\n",
       "    (_input_weights_regularizer): LaplaceL2norm(\n",
       "      (laplace): Laplace()\n",
       "    )\n",
       "    (features): Sequential(\n",
       "      (layer0): Sequential(\n",
       "        (hermite_conv): HermiteConv2D(\n",
       "          (rotate_hermite): RotateHermite(\n",
       "            (Rs): ParameterList(\n",
       "                (0): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (1): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (2): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (3): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (4): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (5): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (6): Parameter containing: [torch.float32 of size 45x45]\n",
       "                (7): Parameter containing: [torch.float32 of size 45x45]\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (norm): RotationEquivariantBatchNorm2D(\n",
       "          (batch_norm): BatchNorm1d(8, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer1): Sequential(\n",
       "        (hermite_conv): HermiteConv2D(\n",
       "          (rotate_hermite): RotateHermite(\n",
       "            (Rs): ParameterList(\n",
       "                (0): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (1): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (2): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (3): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (4): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (5): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (6): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (7): Parameter containing: [torch.float32 of size 28x28]\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (norm): RotationEquivariantBatchNorm2D(\n",
       "          (batch_norm): BatchNorm1d(8, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (hermite_conv): HermiteConv2D(\n",
       "          (rotate_hermite): RotateHermite(\n",
       "            (Rs): ParameterList(\n",
       "                (0): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (1): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (2): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (3): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (4): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (5): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (6): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (7): Parameter containing: [torch.float32 of size 28x28]\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (norm): RotationEquivariantBatchNorm2D(\n",
       "          (batch_norm): BatchNorm1d(8, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (nonlin): AdaptiveELU()\n",
       "      )\n",
       "      (layer3): Sequential(\n",
       "        (hermite_conv): HermiteConv2D(\n",
       "          (rotate_hermite): RotateHermite(\n",
       "            (Rs): ParameterList(\n",
       "                (0): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (1): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (2): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (3): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (4): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (5): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (6): Parameter containing: [torch.float32 of size 28x28]\n",
       "                (7): Parameter containing: [torch.float32 of size 28x28]\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (norm): RotationEquivariantBatchNorm2D(\n",
       "          (batch_norm): BatchNorm1d(8, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  ) [RotationEquivariant2dCore regularizers: gamma_hidden = 500000|gamma_input = 500|skip = 0]\n",
       "  \n",
       "  (readout): MultipleFullGaussian2d(\n",
       "    (dynamic29515-10-12-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (64 x 28 x 56 -> 7863) 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",
       "    (dynamic29755-2-8-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (64 x 28 x 56 -> 8122) 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",
       "    (dynamic29623-4-9-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (64 x 28 x 56 -> 7908) 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",
       "    (dynamic29712-5-9-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (64 x 28 x 56 -> 7939) 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",
       "    (dynamic29647-19-8-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (64 x 28 x 56 -> 8202) 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",
       "    (dynamic29515-10-12-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29755-2-8-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29623-4-9-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29712-5-9-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29647-19-8-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "  (gru_module): GRU_Module(\n",
       "    (gru): ConvGRUCell(\n",
       "      (reset_gate_input): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "      (reset_gate_hidden): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "      (update_gate_input): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "      (update_gate_hidden): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "      (out_gate_input): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "      (out_gate_hidden): Conv2d(64, 64, kernel_size=(9, 9), stride=(1, 1), padding=(4, 4))\n",
       "    ) [ConvGRUCell regularizers: gamma_rec = 0]\n",
       "    \n",
       "  )\n",
       "  (nonlinearity_fn): ELU(alpha=1.0)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gru_2d_model_equivariant"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f113f2-76ae-4e84-8bed-b7b04aabac11",
   "metadata": {},
   "source": [
    "## Factorized Benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dec6eb6f-71c0-468e-b876-c5420f73b48f",
   "metadata": {},
   "outputs": [],
   "source": [
    "factorised_3D_core_dict = dict(\n",
    "    input_channels=3,\n",
    "    hidden_channels=[32, 64, 128],\n",
    "    spatial_input_kernel=(11,11),\n",
    "    temporal_input_kernel=11,\n",
    "    spatial_hidden_kernel=(5,5),\n",
    "    temporal_hidden_kernel=5,\n",
    "    stride=1,\n",
    "    layers=3,\n",
    "    gamma_input_spatial=10,\n",
    "    gamma_input_temporal=0.01, \n",
    "    bias=True, \n",
    "    hidden_nonlinearities='elu', \n",
    "    x_shift=0, \n",
    "    y_shift=0,\n",
    "    batch_norm=True, \n",
    "    laplace_padding=None,\n",
    "    input_regularizer='LaplaceL2norm',\n",
    "    padding=False,\n",
    "    final_nonlin=True,\n",
    "    momentum=0.7\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a4986c98-65b9-4eb7-8cc6-22272b3c3800",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/cores/base.py:82: UserWarning: The batch_norm is applied to all layers\n",
      "  warnings.warn(f\"The {attr} is applied to all layers\", UserWarning)\n",
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/cores/base.py:82: UserWarning: The bias is applied to all layers\n",
      "  warnings.warn(f\"The {attr} is applied to all layers\", UserWarning)\n",
      "/user/turishcheva/anaconda3/envs/sensroium_2023/lib/python3.12/site-packages/neuralpredictors/layers/cores/base.py:82: UserWarning: The batch_norm_scale is applied to all layers\n",
      "  warnings.warn(f\"The {attr} is applied to all layers\", UserWarning)\n"
     ]
    }
   ],
   "source": [
    "factorised_3d_model = make_video_model(\n",
    "    data_loaders,\n",
    "    seed,\n",
    "    core_dict=factorised_3D_core_dict,\n",
    "    core_type='3D_factorised',\n",
    "    readout_dict=readout_dict.copy(),\n",
    "    readout_type='gaussian',               \n",
    "    use_gru=False,\n",
    "    gru_dict=None,\n",
    "    use_shifter=True,\n",
    "    shifter_dict=shifter_dict,\n",
    "    shifter_type='MLP',\n",
    "    deeplake_ds=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1b947b91-ae4a-40f9-9162-c3367ac13517",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VideoFiringRateEncoder(\n",
       "  (core): Factorized3dCore(\n",
       "    (_input_weight_regularizer): LaplaceL2norm(\n",
       "      (laplace): Laplace()\n",
       "    )\n",
       "    (temporal_regularizer): DepthLaplaceL21d(\n",
       "      (laplace): Laplace1d()\n",
       "    )\n",
       "    (features): Sequential(\n",
       "      (layer0): Sequential(\n",
       "        (conv_spatial): Conv3d(3, 32, kernel_size=(1, 11, 11), stride=(1, 1, 1))\n",
       "        (conv_temporal): Conv3d(32, 32, kernel_size=(11, 1, 1), stride=(1, 1, 1))\n",
       "        (norm): BatchNorm3d(32, eps=1e-05, momentum=0.7, affine=True, track_running_stats=True)\n",
       "        (nonlin): ELU(alpha=1.0)\n",
       "      )\n",
       "      (layer1): Sequential(\n",
       "        (conv_spatial_1): Conv3d(32, 64, kernel_size=(1, 5, 5), stride=(1, 1, 1))\n",
       "        (conv_temporal_1): Conv3d(64, 64, kernel_size=(5, 1, 1), stride=(1, 1, 1))\n",
       "        (norm): BatchNorm3d(64, eps=1e-05, momentum=0.7, affine=True, track_running_stats=True)\n",
       "        (nonlin): ELU(alpha=1.0)\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (conv_spatial_2): Conv3d(64, 128, kernel_size=(1, 5, 5), stride=(1, 1, 1))\n",
       "        (conv_temporal_2): Conv3d(128, 128, kernel_size=(5, 1, 1), stride=(1, 1, 1))\n",
       "        (norm): BatchNorm3d(128, eps=1e-05, momentum=0.7, affine=True, track_running_stats=True)\n",
       "        (nonlin): ELU(alpha=1.0)\n",
       "      )\n",
       "    )\n",
       "  ) [Factorized3dCore regularizers: gamma_input_spatial = 10|gamma_input_temporal = 0.01]\n",
       "  \n",
       "  (readout): MultipleFullGaussian2d(\n",
       "    (dynamic29515-10-12-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (128 x 18 x 46 -> 7863) 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",
       "    (dynamic29755-2-8-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (128 x 18 x 46 -> 8122) 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",
       "    (dynamic29623-4-9-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (128 x 18 x 46 -> 7908) 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",
       "    (dynamic29712-5-9-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (128 x 18 x 46 -> 7939) 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",
       "    (dynamic29647-19-8-Video-9b4f6a1a067fe51e15306b9628efea20): full FullGaussian2d (128 x 18 x 46 -> 8202) 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",
       "    (dynamic29515-10-12-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29755-2-8-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29623-4-9-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29712-5-9-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "    (dynamic29647-19-8-Video-9b4f6a1a067fe51e15306b9628efea20): 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",
       "  (nonlinearity_fn): ELU(alpha=1.0)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "factorised_3d_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e7d0831-afae-4cbd-bfa8-34ee5819cd83",
   "metadata": {},
   "source": [
    "## Traning "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5caf0893-9e37-4eea-82ed-52144485bdf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_fn = \"sensorium.training.video_training_loop.standard_trainer\"\n",
    "\n",
    "trainer_config = {\n",
    "    'dataloaders' : data_loaders,\n",
    "    'seed' : seed,\n",
    "    'use_wandb' : False,\n",
    "    'verbose': True,\n",
    "    'lr_decay_steps': 4,\n",
    "    'lr_init': 0.005,\n",
    "    'device' : device,\n",
    "    'detach_core' : False,\n",
    "    # todo - put this to True if you are using deeplake\n",
    "    # first connections to deeplake may take up for 10 mins\n",
    "    'deeplake_ds' : False,\n",
    "    'checkpoint_save_path': 'benchmarks/'\n",
    "                 }\n",
    "\n",
    "trainer = get_trainer(trainer_fn=trainer_fn, \n",
    "                 trainer_config=trainer_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59acc47-d3b3-4607-8916-7af4007dc980",
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace with factorised_3d_model to train factorized benchmark\n",
    "validation_score, trainer_output, state_dict = trainer(gru_2d_model_equivariant)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "607d5e15-5072-43cc-9a96-b6b457e44e17",
   "metadata": {},
   "source": [
    "## Make submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfdd38d3-8e3e-4a2a-b1f8-2ba0a99fb909",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sensorium.utility.submission import generate_submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7955f8-83a8-4500-8707-ab83c06b4918",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_loaders2 = mouse_video_loader(\n",
    "    paths=mice,\n",
    "    batch_size=1,\n",
    "    scale=1,\n",
    "    max_frame=None,\n",
    "    frames=350, # take all frames\n",
    "    offset=-1,\n",
    "    include_behavior=True,\n",
    "    include_pupil_centers=True,\n",
    "    cuda=device!='cpu',\n",
    "    to_cut=False, # take all frames for submission\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4318c2f-f73b-4fcf-86c3-0cb032fc7300",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_submission(data_loaders2, gru_2d_model_equivariant, deeplake_ds=False, device=device, path='benchmarks/')"
   ]
  }
 ],
 "metadata": {
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}