train_cnn_gru_shifter_table_2_sens_orig.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Subset, DataLoader, ConcatDataset\n",
    "from mouse_model.data_utils_new import MouseDatasetSegNewBehav\n",
    "import numpy as np\n",
    "from mouse_model.evaluation import cor_in_time\n",
    "from sklearn.metrics import r2_score, mean_squared_error\n",
    "import random\n",
    "import os\n",
    "import time\n",
    "import torch.nn.init as init\n",
    "from torch.nn import functional as F\n",
    "from kornia.geometry.transform import get_affine_matrix2d, warp_affine\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Shifter(nn.Module):\n",
    "    def __init__(self, input_dim=4, output_dim=3, hidden_dim=256, seq_len=8):\n",
    "        super().__init__()\n",
    "        self.seq_len = seq_len\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            nn.BatchNorm1d(input_dim),\n",
    "            nn.Linear(input_dim, hidden_dim),\n",
    "            nn.BatchNorm1d(hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.BatchNorm1d(hidden_dim),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_dim, output_dim),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.bias = nn.Parameter(torch.zeros(3))\n",
    "    def forward(self, x):\n",
    "        x = x.reshape(-1,self.input_dim )\n",
    "        x = self.layers(x)\n",
    "        x0 = (x[...,0] + self.bias[0]) * 80/5.5\n",
    "        x1 = (x[...,1] + self.bias[1]) * 60/5.5\n",
    "        x2 = (x[...,2] + self.bias[2]) * 180/4\n",
    "        x = torch.stack([x0, x1, x2], dim=-1)\n",
    "        x = x.reshape(-1,self.seq_len,self.output_dim)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrintLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(PrintLayer, self).__init__()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        print(x.shape)\n",
    "        return x\n",
    "    \n",
    "def size_helper(in_length, kernel_size, padding=0, dilation=1, stride=1):\n",
    "    # https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d\n",
    "    res = in_length + 2 * padding - dilation * (kernel_size - 1) - 1\n",
    "    res /= stride\n",
    "    res += 1\n",
    "    return np.floor(res)\n",
    "\n",
    "# CNN, the last fully connected layer maps to output_dim\n",
    "class VisualEncoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, output_dim, input_shape=(60, 80), k1=7, k2=7, k3=7):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.input_shape = (60, 80)\n",
    "        out_shape_0 = size_helper(in_length=input_shape[0], kernel_size=k1, stride=2)\n",
    "        out_shape_0 = size_helper(in_length=out_shape_0, kernel_size=k2, stride=2)\n",
    "        out_shape_0 = size_helper(in_length=out_shape_0, kernel_size=k3, stride=2)\n",
    "        out_shape_1 = size_helper(in_length=input_shape[1], kernel_size=k1, stride=2)\n",
    "        out_shape_1 = size_helper(in_length=out_shape_1, kernel_size=k2, stride=2)\n",
    "        out_shape_1 = size_helper(in_length=out_shape_1, kernel_size=k3, stride=2)\n",
    "        self.output_shape = (int(out_shape_0), int(out_shape_1)) # shape of the final feature map\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=1, out_channels=128, kernel_size=k1, stride=2),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.5),\n",
    "            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=k2, stride=2),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.5),\n",
    "            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=k3, stride=2),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.5),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(480, output_dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.layers(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    \n",
    "# may consider adding an activation after linear\n",
    "class BehavEncoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, behav_dim, output_dim):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            nn.BatchNorm1d(behav_dim),\n",
    "            nn.Linear(behav_dim, output_dim),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.layers(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "class LSTMPerNeuronCombiner(nn.Module):\n",
    "    \n",
    "    def __init__(self, num_neurons, behav_dim, k1, k2, k3, seq_len, hidden_size=512):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.seq_len = seq_len\n",
    "        self.behav_dim = behav_dim\n",
    "        self.num_neurons = num_neurons\n",
    "        self.shifter = Shifter(seq_len = seq_len)\n",
    "        self.visual_encoder = VisualEncoder(output_dim=num_neurons, k1=k1, k2=k2, k3=k3)\n",
    "        self.behav_encoder = BehavEncoder(behav_dim=behav_dim, output_dim=num_neurons)\n",
    "        self.bn = nn.BatchNorm1d(3) # apply bn to vis_feats, beh_feats, prod\n",
    "        self.lstm_net = nn.GRU(input_size=num_neurons*3, hidden_size=hidden_size, num_layers=1, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, num_neurons)\n",
    "        self.softplus = nn.Softplus() # we could also do relu or elu offset by 1\n",
    "        \n",
    "    def forward(self, images, behav):\n",
    "        if args.shifter:\n",
    "            bs = images.size()[0]\n",
    "            behav_shifter = torch.concat((behav[...,4].unsqueeze(-1),   # theta\n",
    "                                          behav[...,3].unsqueeze(-1),   # phi\n",
    "                                          behav[...,1].unsqueeze(-1),  # pitch\n",
    "                                         behav[...,2].unsqueeze(-1),  # roll\n",
    "                                         ), dim=-1)  \n",
    "            shift_param = self.shifter(behav_shifter)  \n",
    "            shift_param = shift_param.reshape(-1,3)\n",
    "            scale_param = torch.ones_like(shift_param[..., 0:2]).to(shift_param.device)\n",
    "            affine_mat = get_affine_matrix2d(\n",
    "                                            translations=shift_param[..., 0:2] ,\n",
    "                                             scale = scale_param, \n",
    "                                             center =torch.repeat_interleave(torch.tensor([[30,40]], dtype=torch.float), \n",
    "                                                                            bs*self.seq_len, dim=0).to(shift_param.device), \n",
    "                                             angle=shift_param[..., 2])\n",
    "            affine_mat = affine_mat[:, :2, :]\n",
    "            images = warp_affine(images.reshape(-1,1,60,80), affine_mat, dsize=(60,80)).reshape(bs, self.seq_len,1,60,80)\n",
    "\n",
    "        if self.behav_dim == 3:\n",
    "            behav = torch.concat((behav[...,0].unsqueeze(-1),  \n",
    "                                behav[...,5].unsqueeze(-1), \n",
    "                                behav[...,-1].unsqueeze(-1),\n",
    "                            ), dim=-1)  \n",
    "        # get visual behavioral features in time\n",
    "        vis_beh_feats = []\n",
    "        for i in range(self.seq_len):\n",
    "            v = self.visual_encoder(images[:, i, :, :, :])\n",
    "            b = self.behav_encoder(behav[:, i, :])\n",
    "            vb = v * b\n",
    "            vis_beh_feat = torch.stack([v, b, vb], axis=1)\n",
    "            vis_beh_feat = self.bn(vis_beh_feat)\n",
    "            vis_beh_feats.append(vis_beh_feat)\n",
    "        vis_beh_feats = torch.stack(vis_beh_feats, axis=1)\n",
    "        \n",
    "        # flatten features to (batch_size, seq_len, num_neurons*3)\n",
    "        vis_beh_feats = torch.flatten(vis_beh_feats, start_dim=2)\n",
    "        \n",
    "        # get LSTM output\n",
    "        output, _ = self.lstm_net(vis_beh_feats)\n",
    "        output = output[:, -1, :] # extract the last hidden state\n",
    "        \n",
    "        # fully connected layer and activation function\n",
    "        output = self.fc(output)\n",
    "        pred_spikes = self.softplus(output)\n",
    "\n",
    "        return pred_spikes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "class Args:\n",
    "    \n",
    "    seed = 0\n",
    "    file_id = None\n",
    "    epochs = 50\n",
    "    batch_size = 256\n",
    "    learning_rate = 0.0002\n",
    "    l1_reg_w = 1\n",
    "    seq_len = None\n",
    "    num_neurons = None\n",
    "    behav_mode = None\n",
    "    behav_dim = None\n",
    "    best_val_path = None\n",
    "    best_train_path = None\n",
    "    vid_type = \"vid_mean\"\n",
    "    segment_num = 10\n",
    "    hidden_size = 512\n",
    "    shifter = True\n",
    "    \n",
    "args=Args()\n",
    "\n",
    "seed = args.seed\n",
    "random.seed(seed)\n",
    "os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_train_val_ds():\n",
    "    ds_list = [MouseDatasetSegNewBehav(file_id=args.file_id, segment_num=args.segment_num, seg_idx=i, data_split=\"train\", \n",
    "                               vid_type=args.vid_type, seq_len=args.seq_len, predict_offset=1, \n",
    "                                       behav_mode=args.behav_mode, norm_mode=\"01\") \n",
    "               for i in range(args.segment_num)]\n",
    "    train_ds, val_ds = [], []\n",
    "    for ds in ds_list:\n",
    "        train_ratio = 0.8\n",
    "        train_ds_len = int(len(ds) * train_ratio)\n",
    "        train_ds.append(Subset(ds, np.arange(0, train_ds_len, 1)))\n",
    "        val_ds.append(Subset(ds, np.arange(train_ds_len, len(ds), 1)))\n",
    "    train_ds = ConcatDataset(train_ds)\n",
    "    val_ds = ConcatDataset(val_ds)\n",
    "    print(len(train_ds), len(val_ds))\n",
    "    return train_ds, val_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_test_ds():\n",
    "    test_ds = [MouseDatasetSegNewBehav(file_id=args.file_id, segment_num=args.segment_num, seg_idx=i, data_split=\"test\", \n",
    "                               vid_type=args.vid_type, seq_len=args.seq_len, predict_offset=1, \n",
    "                                       behav_mode=args.behav_mode, norm_mode=\"01\") \n",
    "               for i in range(args.segment_num)]\n",
    "    test_ds = ConcatDataset(test_ds)\n",
    "    return test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model():\n",
    "    \n",
    "    torch.manual_seed(args.seed)\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.manual_seed(args.seed)\n",
    "    \n",
    "    train_ds, val_ds = load_train_val_ds()\n",
    "\n",
    "    train_dataloader = DataLoader(dataset=train_ds, batch_size=args.batch_size, shuffle=True, num_workers=8)\n",
    "    val_dataloader = DataLoader(dataset=val_ds, batch_size=args.batch_size, shuffle=False, num_workers=8)\n",
    "    \n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)\n",
    "\n",
    "    best_train_spike_loss = np.inf\n",
    "    best_val_spike_loss = np.inf\n",
    "    train_loss_list = []\n",
    "    val_loss_list = []\n",
    "\n",
    "    # start training\n",
    "    ct = 0\n",
    "    \n",
    "    for epoch in range(args.epochs):\n",
    "\n",
    "        print(\"Start epoch\", epoch)\n",
    "\n",
    "        model.train()\n",
    "\n",
    "        epoch_train_loss, epoch_train_spike_loss = 0, 0\n",
    "\n",
    "        for (image, behav, spikes) in train_dataloader:\n",
    "\n",
    "            image, behav, spikes = image.to(device), behav.to(device), spikes.to(device)\n",
    "            \n",
    "            pred = model(image, behav)\n",
    "\n",
    "            spike_loss = nn.functional.poisson_nll_loss(pred, spikes, reduction='mean', log_input=False)\n",
    "            \n",
    "            l1_reg, l1_reg_num_param = 0.0, 0\n",
    "            for name, param in model.named_parameters():\n",
    "                if name == \"behav_encoder.layers.1.weight\":\n",
    "                    l1_reg += param.abs().sum()\n",
    "                    l1_reg_num_param += param.shape[0]*param.shape[1]\n",
    "            l1_reg /= l1_reg_num_param\n",
    "\n",
    "            total_loss = spike_loss + args.l1_reg_w * l1_reg\n",
    "\n",
    "            epoch_train_loss += total_loss.item()\n",
    "            epoch_train_spike_loss += spike_loss.item()\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            total_loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        epoch_train_loss = epoch_train_loss / len(train_dataloader)\n",
    "        epoch_train_spike_loss = epoch_train_spike_loss / len(train_dataloader)\n",
    "\n",
    "        train_loss_list.append(epoch_train_loss)\n",
    "        \n",
    "        print(\"Epoch {} train loss: {}\".format(epoch, epoch_train_loss))\n",
    "\n",
    "        if epoch_train_spike_loss < best_train_spike_loss:\n",
    "\n",
    "            print(\"save train model at epoch\", epoch)\n",
    "            torch.save(model.state_dict(), args.best_train_path)\n",
    "            best_train_spike_loss = epoch_train_spike_loss\n",
    "\n",
    "        model.eval()\n",
    "\n",
    "        epoch_val_spike_loss = 0\n",
    "\n",
    "        with torch.no_grad():      \n",
    "\n",
    "            for (image, behav, spikes) in val_dataloader:\n",
    "\n",
    "                image, behav, spikes = image.to(device), behav.to(device), spikes.to(device)\n",
    "\n",
    "                pred = model(image, behav)\n",
    "\n",
    "                loss = nn.functional.poisson_nll_loss(pred, spikes, reduction='mean', log_input=False)\n",
    "\n",
    "                epoch_val_spike_loss += loss.item()\n",
    "\n",
    "        epoch_val_spike_loss = epoch_val_spike_loss / len(val_dataloader)\n",
    "\n",
    "        val_loss_list.append(epoch_val_spike_loss)\n",
    "        \n",
    "        print(\"Epoch {} val loss: {}\".format(epoch, epoch_val_spike_loss))\n",
    "        \n",
    "        if epoch_val_spike_loss < best_val_spike_loss:\n",
    "            ct = 0\n",
    "\n",
    "            print(\"save val model at epoch\", epoch)\n",
    "            torch.save(model.state_dict(), args.best_val_path)\n",
    "            best_val_spike_loss = epoch_val_spike_loss\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct >=5:\n",
    "                print('stop training')\n",
    "                break\n",
    "\n",
    "        print(\"End epoch\", epoch)\n",
    "        \n",
    "    return train_loss_list, val_loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT all 1\n",
      "30120 7540\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 0.9536194478051138\n",
      "save train model at epoch 0\n",
      "Epoch 0 val loss: 0.657522960503896\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.8925990654250323\n",
      "save train model at epoch 1\n",
      "Epoch 1 val loss: 0.6512650827566783\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.8599661346209251\n",
      "save train model at epoch 2\n",
      "Epoch 2 val loss: 0.639069265127182\n",
      "save val model at epoch 2\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.8299451116788186\n",
      "save train model at epoch 3\n",
      "Epoch 3 val loss: 0.6344329416751862\n",
      "save val model at epoch 3\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.8037741032697386\n",
      "save train model at epoch 4\n",
      "Epoch 4 val loss: 0.6259421209494272\n",
      "save val model at epoch 4\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.7812813040563615\n",
      "save train model at epoch 5\n",
      "Epoch 5 val loss: 0.6260844190915426\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.760532091734773\n",
      "save train model at epoch 6\n",
      "Epoch 6 val loss: 0.6195620973904927\n",
      "save val model at epoch 6\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.7416726613448839\n",
      "save train model at epoch 7\n",
      "Epoch 7 val loss: 0.6246476928393047\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "Epoch 8 train loss: 0.7239675380415835\n",
      "save train model at epoch 8\n",
      "Epoch 8 val loss: 0.618780501683553\n",
      "save val model at epoch 8\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "Epoch 9 train loss: 0.7066998405981872\n",
      "save train model at epoch 9\n",
      "Epoch 9 val loss: 0.6140351553757986\n",
      "save val model at epoch 9\n",
      "End epoch 9\n",
      "Start epoch 10\n",
      "Epoch 10 train loss: 0.6906642060158616\n",
      "save train model at epoch 10\n",
      "Epoch 10 val loss: 0.6121026297410329\n",
      "save val model at epoch 10\n",
      "End epoch 10\n",
      "Start epoch 11\n",
      "Epoch 11 train loss: 0.6759406510045973\n",
      "save train model at epoch 11\n",
      "Epoch 11 val loss: 0.616668297847112\n",
      "End epoch 11\n",
      "Start epoch 12\n",
      "Epoch 12 train loss: 0.6624602331953534\n",
      "save train model at epoch 12\n",
      "Epoch 12 val loss: 0.6134611030419668\n",
      "End epoch 12\n",
      "Start epoch 13\n",
      "Epoch 13 train loss: 0.6503296681379868\n",
      "save train model at epoch 13\n",
      "Epoch 13 val loss: 0.61272714138031\n",
      "End epoch 13\n",
      "Start epoch 14\n",
      "Epoch 14 train loss: 0.6396056873313452\n",
      "save train model at epoch 14\n",
      "Epoch 14 val loss: 0.6085595945517223\n",
      "save val model at epoch 14\n",
      "End epoch 14\n",
      "Start epoch 15\n",
      "Epoch 15 train loss: 0.6303105985714217\n",
      "save train model at epoch 15\n",
      "Epoch 15 val loss: 0.6071246027946472\n",
      "save val model at epoch 15\n",
      "End epoch 15\n",
      "Start epoch 16\n",
      "Epoch 16 train loss: 0.6224249030573893\n",
      "save train model at epoch 16\n",
      "Epoch 16 val loss: 0.6065837343533834\n",
      "save val model at epoch 16\n",
      "End epoch 16\n",
      "Start epoch 17\n",
      "Epoch 17 train loss: 0.615689852480161\n",
      "save train model at epoch 17\n",
      "Epoch 17 val loss: 0.6064306954542796\n",
      "save val model at epoch 17\n",
      "End epoch 17\n",
      "Start epoch 18\n",
      "Epoch 18 train loss: 0.6102586822994684\n",
      "save train model at epoch 18\n",
      "Epoch 18 val loss: 0.6048961738745372\n",
      "save val model at epoch 18\n",
      "End epoch 18\n",
      "Start epoch 19\n",
      "Epoch 19 train loss: 0.6054722984968606\n",
      "save train model at epoch 19\n",
      "Epoch 19 val loss: 0.608334356546402\n",
      "End epoch 19\n",
      "Start epoch 20\n",
      "Epoch 20 train loss: 0.6012983322143555\n",
      "save train model at epoch 20\n",
      "Epoch 20 val loss: 0.605909385283788\n",
      "End epoch 20\n",
      "Start epoch 21\n",
      "Epoch 21 train loss: 0.598170607271841\n",
      "save train model at epoch 21\n",
      "Epoch 21 val loss: 0.6050628821055094\n",
      "End epoch 21\n",
      "Start epoch 22\n",
      "Epoch 22 train loss: 0.5957848434731111\n",
      "save train model at epoch 22\n",
      "Epoch 22 val loss: 0.6051699002583821\n",
      "End epoch 22\n",
      "Start epoch 23\n",
      "Epoch 23 train loss: 0.5938358054322711\n",
      "save train model at epoch 23\n",
      "Epoch 23 val loss: 0.6089443862438202\n",
      "stop training\n",
      "101521_J559NC all 1\n",
      "42410 10610\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 1.0222428204065346\n",
      "save train model at epoch 0\n",
      "Epoch 0 val loss: 0.7233363503501529\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.9558232782116856\n",
      "save train model at epoch 1\n",
      "Epoch 1 val loss: 0.7110971964540935\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.9137356306415007\n",
      "save train model at epoch 2\n",
      "Epoch 2 val loss: 0.7067791776997703\n",
      "save val model at epoch 2\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.8796412169215191\n",
      "save train model at epoch 3\n",
      "Epoch 3 val loss: 0.7022604913938613\n",
      "save val model at epoch 3\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.8490831241550216\n",
      "save train model at epoch 4\n",
      "Epoch 4 val loss: 0.7158386082876296\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.8224388913935926\n",
      "save train model at epoch 5\n",
      "Epoch 5 val loss: 0.699059084767387\n",
      "save val model at epoch 5\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.7980093722601971\n",
      "save train model at epoch 6\n",
      "Epoch 6 val loss: 0.6979751047633943\n",
      "save val model at epoch 6\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.7774155107607325\n",
      "save train model at epoch 7\n",
      "Epoch 7 val loss: 0.6989460431394123\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "Epoch 8 train loss: 0.7602535705250429\n",
      "save train model at epoch 8\n",
      "Epoch 8 val loss: 0.7284048029354641\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "Epoch 9 train loss: 0.7438621032668884\n",
      "save train model at epoch 9\n",
      "Epoch 9 val loss: 0.6981930306979588\n",
      "End epoch 9\n",
      "Start epoch 10\n",
      "Epoch 10 train loss: 0.7288675308227539\n",
      "save train model at epoch 10\n",
      "Epoch 10 val loss: 0.7223884434927077\n",
      "End epoch 10\n",
      "Start epoch 11\n",
      "Epoch 11 train loss: 0.7160949434142515\n",
      "save train model at epoch 11\n",
      "Epoch 11 val loss: 0.695103318918319\n",
      "save val model at epoch 11\n",
      "End epoch 11\n",
      "Start epoch 12\n",
      "Epoch 12 train loss: 0.7048434586410063\n",
      "save train model at epoch 12\n",
      "Epoch 12 val loss: 0.6959465486662728\n",
      "End epoch 12\n",
      "Start epoch 13\n",
      "Epoch 13 train loss: 0.6963931562670742\n",
      "save train model at epoch 13\n",
      "Epoch 13 val loss: 0.6949405556633359\n",
      "save val model at epoch 13\n",
      "End epoch 13\n",
      "Start epoch 14\n",
      "Epoch 14 train loss: 0.6898520807903933\n",
      "save train model at epoch 14\n",
      "Epoch 14 val loss: 0.6972257849716005\n",
      "End epoch 14\n",
      "Start epoch 15\n",
      "Epoch 15 train loss: 0.6847932805497963\n",
      "save train model at epoch 15\n",
      "Epoch 15 val loss: 0.6983260072412945\n",
      "End epoch 15\n",
      "Start epoch 16\n",
      "Epoch 16 train loss: 0.6810111073126276\n",
      "save train model at epoch 16\n",
      "Epoch 16 val loss: 0.6917941996029445\n",
      "save val model at epoch 16\n",
      "End epoch 16\n",
      "Start epoch 17\n",
      "Epoch 17 train loss: 0.6786896212991461\n",
      "save train model at epoch 17\n",
      "Epoch 17 val loss: 0.69880968048459\n",
      "End epoch 17\n",
      "Start epoch 18\n",
      "Epoch 18 train loss: 0.6768268817160503\n",
      "save train model at epoch 18\n",
      "Epoch 18 val loss: 0.6919687177453723\n",
      "End epoch 18\n",
      "Start epoch 19\n",
      "Epoch 19 train loss: 0.6750685306916754\n",
      "save train model at epoch 19\n",
      "Epoch 19 val loss: 0.690596661397389\n",
      "save val model at epoch 19\n",
      "End epoch 19\n",
      "Start epoch 20\n",
      "Epoch 20 train loss: 0.6740130419472614\n",
      "save train model at epoch 20\n",
      "Epoch 20 val loss: 0.6884406237375169\n",
      "save val model at epoch 20\n",
      "End epoch 20\n",
      "Start epoch 21\n",
      "Epoch 21 train loss: 0.6735428219818207\n",
      "save train model at epoch 21\n",
      "Epoch 21 val loss: 0.692036197299049\n",
      "End epoch 21\n",
      "Start epoch 22\n",
      "Epoch 22 train loss: 0.6719337691025562\n",
      "save train model at epoch 22\n",
      "Epoch 22 val loss: 0.688139830316816\n",
      "save val model at epoch 22\n",
      "End epoch 22\n",
      "Start epoch 23\n",
      "Epoch 23 train loss: 0.6715101987482553\n",
      "save train model at epoch 23\n",
      "Epoch 23 val loss: 0.6985566502525693\n",
      "End epoch 23\n",
      "Start epoch 24\n",
      "Epoch 24 train loss: 0.6704899298857494\n",
      "save train model at epoch 24\n",
      "Epoch 24 val loss: 0.6918258056754157\n",
      "End epoch 24\n",
      "Start epoch 25\n",
      "Epoch 25 train loss: 0.6694481128669647\n",
      "save train model at epoch 25\n",
      "Epoch 25 val loss: 0.6895291109879812\n",
      "End epoch 25\n",
      "Start epoch 26\n",
      "Epoch 26 train loss: 0.6691093814660267\n",
      "save train model at epoch 26\n",
      "Epoch 26 val loss: 0.6874276640869322\n",
      "save val model at epoch 26\n",
      "End epoch 26\n",
      "Start epoch 27\n",
      "Epoch 27 train loss: 0.6677648622587503\n",
      "save train model at epoch 27\n",
      "Epoch 27 val loss: 0.7014405670620146\n",
      "End epoch 27\n",
      "Start epoch 28\n",
      "Epoch 28 train loss: 0.6673761128661144\n",
      "save train model at epoch 28\n",
      "Epoch 28 val loss: 0.688664266041347\n",
      "End epoch 28\n",
      "Start epoch 29\n",
      "Epoch 29 train loss: 0.6664424890495209\n",
      "save train model at epoch 29\n",
      "Epoch 29 val loss: 0.6898993906520662\n",
      "End epoch 29\n",
      "Start epoch 30\n",
      "Epoch 30 train loss: 0.6660166038088051\n",
      "save train model at epoch 30\n",
      "Epoch 30 val loss: 0.6861844602085295\n",
      "save val model at epoch 30\n",
      "End epoch 30\n",
      "Start epoch 31\n",
      "Epoch 31 train loss: 0.6654346613998873\n",
      "save train model at epoch 31\n",
      "Epoch 31 val loss: 0.6964799194108873\n",
      "End epoch 31\n",
      "Start epoch 32\n",
      "Epoch 32 train loss: 0.6649380854095321\n",
      "save train model at epoch 32\n",
      "Epoch 32 val loss: 0.6906928306534177\n",
      "End epoch 32\n",
      "Start epoch 33\n",
      "Epoch 33 train loss: 0.6642336698181658\n",
      "save train model at epoch 33\n",
      "Epoch 33 val loss: 0.6911880998384385\n",
      "End epoch 33\n",
      "Start epoch 34\n",
      "Epoch 34 train loss: 0.6633917491838156\n",
      "save train model at epoch 34\n",
      "Epoch 34 val loss: 0.6849851551510039\n",
      "save val model at epoch 34\n",
      "End epoch 34\n",
      "Start epoch 35\n",
      "Epoch 35 train loss: 0.6625706255435944\n",
      "save train model at epoch 35\n",
      "Epoch 35 val loss: 0.6884994648751759\n",
      "End epoch 35\n",
      "Start epoch 36\n",
      "Epoch 36 train loss: 0.6620628226952381\n",
      "save train model at epoch 36\n",
      "Epoch 36 val loss: 0.6861972979136876\n",
      "End epoch 36\n",
      "Start epoch 37\n",
      "Epoch 37 train loss: 0.661542006644858\n",
      "save train model at epoch 37\n",
      "Epoch 37 val loss: 0.686276056936809\n",
      "End epoch 37\n",
      "Start epoch 38\n",
      "Epoch 38 train loss: 0.6607936516583685\n",
      "save train model at epoch 38\n",
      "Epoch 38 val loss: 0.6884887147517431\n",
      "End epoch 38\n",
      "Start epoch 39\n",
      "Epoch 39 train loss: 0.6606108984315252\n",
      "save train model at epoch 39\n",
      "Epoch 39 val loss: 0.6989812808377402\n",
      "stop training\n",
      "110421_J569LT all 1\n",
      "32940 8240\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 1.0505088687867157\n",
      "save train model at epoch 0\n",
      "Epoch 0 val loss: 0.7554999553796017\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.9890705758287001\n",
      "save train model at epoch 1\n",
      "Epoch 1 val loss: 0.7453567710789767\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.954126164432644\n",
      "save train model at epoch 2\n",
      "Epoch 2 val loss: 0.7463454206784567\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.9269472355990447\n",
      "save train model at epoch 3\n",
      "Epoch 3 val loss: 0.7458240787188212\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.9024728384128836\n",
      "save train model at epoch 4\n",
      "Epoch 4 val loss: 0.734693725903829\n",
      "save val model at epoch 4\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.8797297316004139\n",
      "save train model at epoch 5\n",
      "Epoch 5 val loss: 0.7453845323938312\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.8587900852972223\n",
      "save train model at epoch 6\n",
      "Epoch 6 val loss: 0.7322530096227472\n",
      "save val model at epoch 6\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.8393873318221218\n",
      "save train model at epoch 7\n",
      "Epoch 7 val loss: 0.7342896298928694\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "Epoch 8 train loss: 0.8214907184127689\n",
      "save train model at epoch 8\n",
      "Epoch 8 val loss: 0.7326589551838961\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "Epoch 9 train loss: 0.804538677367129\n",
      "save train model at epoch 9\n",
      "Epoch 9 val loss: 0.7377317927100442\n",
      "End epoch 9\n",
      "Start epoch 10\n",
      "Epoch 10 train loss: 0.789181158524151\n",
      "save train model at epoch 10\n",
      "Epoch 10 val loss: 0.7298182122635118\n",
      "save val model at epoch 10\n",
      "End epoch 10\n",
      "Start epoch 11\n",
      "Epoch 11 train loss: 0.7756689130797867\n",
      "save train model at epoch 11\n",
      "Epoch 11 val loss: 0.7321761196309869\n",
      "End epoch 11\n",
      "Start epoch 12\n",
      "Epoch 12 train loss: 0.7627533900645352\n",
      "save train model at epoch 12\n",
      "Epoch 12 val loss: 0.7285894318060442\n",
      "save val model at epoch 12\n",
      "End epoch 12\n",
      "Start epoch 13\n",
      "Epoch 13 train loss: 0.7516327555789504\n",
      "save train model at epoch 13\n",
      "Epoch 13 val loss: 0.7325661092093496\n",
      "End epoch 13\n",
      "Start epoch 14\n",
      "Epoch 14 train loss: 0.7410631128983904\n",
      "save train model at epoch 14\n",
      "Epoch 14 val loss: 0.7326084194761334\n",
      "End epoch 14\n",
      "Start epoch 15\n",
      "Epoch 15 train loss: 0.7321964225103689\n",
      "save train model at epoch 15\n",
      "Epoch 15 val loss: 0.7305446285190005\n",
      "End epoch 15\n",
      "Start epoch 16\n",
      "Epoch 16 train loss: 0.7244846626769664\n",
      "save train model at epoch 16\n",
      "Epoch 16 val loss: 0.7273458350788463\n",
      "save val model at epoch 16\n",
      "End epoch 16\n",
      "Start epoch 17\n",
      "Epoch 17 train loss: 0.7168880183567372\n",
      "save train model at epoch 17\n",
      "Epoch 17 val loss: 0.7336248123284542\n",
      "End epoch 17\n",
      "Start epoch 18\n",
      "Epoch 18 train loss: 0.7105831928955492\n",
      "save train model at epoch 18\n",
      "Epoch 18 val loss: 0.7282048987619805\n",
      "End epoch 18\n",
      "Start epoch 19\n",
      "Epoch 19 train loss: 0.7052290707595589\n",
      "save train model at epoch 19\n",
      "Epoch 19 val loss: 0.7284236532269102\n",
      "End epoch 19\n",
      "Start epoch 20\n",
      "Epoch 20 train loss: 0.7009475965832555\n",
      "save train model at epoch 20\n",
      "Epoch 20 val loss: 0.7340283592542013\n",
      "End epoch 20\n",
      "Start epoch 21\n",
      "Epoch 21 train loss: 0.6984819483387378\n",
      "Epoch 21 val loss: 0.7357475558916727\n",
      "stop training\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')    \n",
    "\n",
    "for file_id, num_neurons in  [(\"070921_J553RT\", 68), (\"101521_J559NC\", 49) , (\"110421_J569LT\", 32)]:\n",
    "    for behav_mode, behav_dim in [(\"all\", 3)]:  \n",
    "    # pass all behav variables\n",
    "    # take 4 (head&eye movement) for shifter\n",
    "    # take 3 (the sensorium original) variable as model main component's input\n",
    "        for seq_len in range(1, 2): \n",
    "            print(file_id, behav_mode, seq_len)\n",
    "            \n",
    "            args.file_id = file_id\n",
    "            args.vid_type = \"vid_mean\"\n",
    "            args.num_neurons = num_neurons\n",
    "            args.shifter=True\n",
    "\n",
    "            args.behav_mode = behav_mode\n",
    "            args.behav_dim = behav_dim\n",
    "            \n",
    "            args.seq_len = seq_len\n",
    "\n",
    "            args.best_train_path = \"/hdd/yuchen/train_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id,  'sens_orig', args.seq_len)\n",
    "            args.best_val_path = \"/hdd/yuchen/val_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id, 'sens_orig' , args.seq_len)\n",
    "\n",
    "            model = LSTMPerNeuronCombiner(num_neurons=args.num_neurons, \n",
    "                                          behav_dim=args.behav_dim, \n",
    "                                          k1=7, k2=7, k3=7, \n",
    "                                          seq_len=args.seq_len,\n",
    "                                          hidden_size=args.hidden_size).to(device)\n",
    "\n",
    "            train_loss_list, val_loss_list = train_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default is smoothing with 2 second, 48 ms per frame\n",
    "def smoothing_with_np_conv(nsp, size=int(2000/48)):\n",
    "    np_conv_res = []\n",
    "    for i in range(nsp.shape[1]):\n",
    "        np_conv_res.append(np.convolve(nsp[:, i], np.ones(size)/size, mode=\"same\"))        \n",
    "    np_conv_res = np.transpose(np.array(np_conv_res))\n",
    "    return np_conv_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, weights_path, dataset, device):\n",
    "\n",
    "    dl = DataLoader(dataset=dataset, batch_size=256, shuffle=False, num_workers=4)\n",
    "    \n",
    "    model.load_state_dict(torch.load(weights_path))\n",
    "\n",
    "    ground_truth_all = []\n",
    "    pred_all = []\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():      \n",
    "        \n",
    "        for (image, behav, spikes) in dl:\n",
    "            \n",
    "            image = image.to(device)\n",
    "            behav = behav.to(device)\n",
    "            \n",
    "            pred = model(image, behav)\n",
    "            \n",
    "            ground_truth_all.append(spikes.numpy())\n",
    "            pred_all.append(pred.cpu().numpy())\n",
    "    \n",
    "    return np.concatenate(pred_all, axis=0), np.concatenate(ground_truth_all, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT all 1\n",
      "30120 7540\n",
      "MSE 0.055057\n",
      "mean corr, 0.623+-0.142\n",
      "101521_J559NC all 1\n",
      "42410 10610\n",
      "MSE 0.082836\n",
      "mean corr, 0.579+-0.135\n",
      "110421_J569LT all 1\n",
      "32940 8240\n",
      "MSE 0.091109\n",
      "mean corr, 0.498+-0.154\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')    \n",
    "\n",
    "# for file_id, num_neurons in [(\"070921_J553RT\", 68)]:\n",
    "for file_id, num_neurons in  [(\"070921_J553RT\", 68), (\"101521_J559NC\", 49) , (\"110421_J569LT\", 32)]:\n",
    "    \n",
    "\n",
    "    for behav_mode, behav_dim in [(\"all\", 3)]:  \n",
    "    # pass all behav variables\n",
    "    # take 4 (head&eye movement) for shifter\n",
    "    # take 3 (the sensorium original) variable as model main component's input\n",
    "\n",
    "        for seq_len in range(1, 2): \n",
    "\n",
    "            print(file_id, behav_mode, seq_len)\n",
    "            \n",
    "            args.file_id = file_id\n",
    "            args.num_neurons = num_neurons\n",
    "\n",
    "            args.behav_mode = behav_mode\n",
    "            args.behav_dim = behav_dim\n",
    "            \n",
    "            args.seq_len = seq_len\n",
    "\n",
    "            args.best_train_path = \"/hdd/yuchen/train_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id,  'sens_orig', args.seq_len)\n",
    "            args.best_val_path = \"/hdd/yuchen/val_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id, 'sens_orig' , args.seq_len)\n",
    "\n",
    "            model = LSTMPerNeuronCombiner(num_neurons=args.num_neurons, \n",
    "                                          behav_dim=args.behav_dim, \n",
    "                                          k1=7, k2=7, k3=7, \n",
    "                                          seq_len=args.seq_len,\n",
    "                                          hidden_size=args.hidden_size).to(device)\n",
    "            \n",
    "            train_ds, val_ds = load_train_val_ds()\n",
    "            test_ds = load_test_ds()\n",
    "            \n",
    "            pred, label = evaluate_model(model, weights_path=args.best_val_path, dataset=test_ds, device=device)\n",
    "            cor_array = cor_in_time(pred, label)\n",
    "        #     print(\"best val model on test dataset, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "            pred = smoothing_with_np_conv(pred)\n",
    "            label = smoothing_with_np_conv(label)\n",
    "            # print(\"R2\", \"{:.6f}\".format(r2_score(label.T, pred.T)))\n",
    "            print(\"MSE\", \"{:.6f}\".format(mean_squared_error(label, pred)))\n",
    "            cor_array = cor_in_time(pred, label)\n",
    "            print(\"mean corr, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "            # print(\"max corr\", \"{:.6f}\".format(np.max(cor_array)))\n",
    "            # print(\"min corr\", \"{:.6f}\".format(np.min(cor_array)))"
   ]
  }
 ],
 "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}