train_cnn_gru_shifter_table_2.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": 3,
   "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.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",
    "        # 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": 4,
   "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": 5,
   "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": 6,
   "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": 7,
   "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": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT sens_orig 1\n",
      "30120 7540\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 0.9059521045725224\n",
      "save train model at epoch 0\n",
      "Epoch 0 val loss: 0.662943039337794\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.8439239308995715\n",
      "save train model at epoch 1\n",
      "Epoch 1 val loss: 0.6534544150034587\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 31\u001b[0m\n\u001b[1;32m     22\u001b[0m args\u001b[38;5;241m.\u001b[39mbest_val_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/hdd/yuchen/val_baseline_\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m_\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m_seq_\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m     23\u001b[0m     args\u001b[38;5;241m.\u001b[39mfile_id, args\u001b[38;5;241m.\u001b[39mbehav_mode, args\u001b[38;5;241m.\u001b[39mseq_len)\n\u001b[1;32m     25\u001b[0m model \u001b[38;5;241m=\u001b[39m LSTMPerNeuronCombiner(num_neurons\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mnum_neurons, \n\u001b[1;32m     26\u001b[0m                               behav_dim\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mbehav_dim, \n\u001b[1;32m     27\u001b[0m                               k1\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m7\u001b[39m, k2\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m7\u001b[39m, k3\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m7\u001b[39m, \n\u001b[1;32m     28\u001b[0m                               seq_len\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mseq_len,\n\u001b[1;32m     29\u001b[0m                               hidden_size\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mhidden_size)\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m---> 31\u001b[0m train_loss_list, val_loss_list \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[7], line 32\u001b[0m, in \u001b[0;36mtrain_model\u001b[0;34m()\u001b[0m\n\u001b[1;32m     28\u001b[0m epoch_train_loss, epoch_train_spike_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m (image, behav, spikes) \u001b[38;5;129;01min\u001b[39;00m train_dataloader:\n\u001b[0;32m---> 32\u001b[0m     image, behav, spikes \u001b[38;5;241m=\u001b[39m \u001b[43mimage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m, behav\u001b[38;5;241m.\u001b[39mto(device), spikes\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     34\u001b[0m     pred \u001b[38;5;241m=\u001b[39m model(image, behav)\n\u001b[1;32m     36\u001b[0m     spike_loss \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mfunctional\u001b[38;5;241m.\u001b[39mpoisson_nll_loss(pred, spikes, reduction\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmean\u001b[39m\u001b[38;5;124m'\u001b[39m, log_input\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')    \n",
    "\n",
    "for file_id, num_neurons in  [(\"070921_J553RT\", 68), (\"110421_J569LT\", 32), (\"101521_J559NC\", 49) ]:\n",
    "    # for behav_mode, behav_dim in [(\"orig_prod\", 21)]:\n",
    "    for behav_mode, behav_dim in [(\"orig\", 6), (\"velo\", 6), (\"all\", 11), (\"orig_prod\", 21), (\"velo_prod\", 21)]:\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, args.behav_mode, args.seq_len)\n",
    "            args.best_val_path = \"/hdd/yuchen/val_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id, args.behav_mode, 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": "markdown",
   "metadata": {},
   "source": [
    "eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT orig_prod 1\n",
      "30120 7540\n",
      "MSE 0.055455\n",
      "mean corr, 0.639+-0.141\n",
      "101521_J559NC orig_prod 1\n",
      "42410 10610\n",
      "MSE 0.081183\n",
      "mean corr, 0.593+-0.132\n",
      "110421_J569LT orig_prod 1\n",
      "32940 8240\n",
      "MSE 0.092540\n",
      "mean corr, 0.484+-0.165\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), (\"110421_J569LT\", 32), (\"101521_J559NC\", 49)]:\n",
    "    \n",
    "\n",
    "    for behav_mode, behav_dim in [(\"orig_prod\", 21)]:\n",
    "    # for behav_mode, behav_dim in [(\"orig\", 6), (\"all\", 11), (\"orig_prod\", 21)]:\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, args.behav_mode, args.seq_len)\n",
    "            args.best_val_path = \"/hdd/yuchen/val_baseline_{}_{}_seq_{}.pth\".format(\n",
    "                args.file_id, args.behav_mode, 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
}