train_efficientnet_shifter_table_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared to autoencoder, this architecture does not include reconstruction.\n",
    "\n",
    "variants to try:\n",
    "* data segment number: 3, 5, 8, 10\n",
    "* eye video vs. head video\n",
    "\n",
    "fixed values:\n",
    "* kernel sizes (3, 3, 3). If we have more time, could try (3, 5, 9) (best config from hyperparameter search)"
   ]
  },
  {
   "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 torchvision import models\n",
    "from sklearn.metrics import r2_score, mean_squared_error\n",
    "import timm\n",
    "from kornia.geometry.transform import get_affine_matrix2d, warp_affine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# useful for printing in nn.Sequential\n",
    "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=3, k2=3, k3=3):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.input_shape = (60, 80)\n",
    "        \n",
    "        efficientnet = timm.create_model('efficientnet_b0', pretrained=True)\n",
    "        efficientnet.conv_stem = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
    "        efficientnet.classifier = nn.Linear(in_features=1280, out_features=output_dim, bias=True)\n",
    "        self.layers =  efficientnet\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.layers(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class Shifter(nn.Module):\n",
    "    def __init__(self, input_dim=4, output_dim=3, hidden_dim=256):\n",
    "        super().__init__()\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/4\n",
    "        x1 = (x[...,1] + self.bias[1]) * 60/4\n",
    "        x2 = (x[...,2] + self.bias[2]) * 180/4\n",
    "        x = torch.stack([x0, x1, x2], dim=-1)\n",
    "        x = x.reshape(-1,1,self.output_dim)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Predictor(nn.Module):\n",
    "    \n",
    "    def __init__(self, num_neurons, k1, k2, k3):\n",
    "\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = VisualEncoder(output_dim=num_neurons, k1=k1, k2=k2, k3=k3)\n",
    "        self.softplus = nn.Softplus()\n",
    "        self.shifter = Shifter()\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(translations=shift_param[..., 0:2] ,\n",
    "                                                scale = scale_param, \n",
    "                                                 center =torch.repeat_interleave(torch.tensor([[30,40]], dtype=torch.float), \n",
    "                                                                                bs*1, dim=0).to(shift_param.device), \n",
    "                                                 angle=shift_param[..., 2])\n",
    "            affine_mat = affine_mat[:, :2, :]\n",
    "            images = warp_affine(images, affine_mat, dsize=(60,80))\n",
    "        \n",
    "        pred = self.encoder(images)\n",
    "        pred = self.softplus(pred)\n",
    "        \n",
    "        return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    \n",
    "    seed = 0\n",
    "    file_id = None\n",
    "    epochs = 100\n",
    "    batch_size = 256\n",
    "    seq_len = 1\n",
    "    num_neurons = None\n",
    "    learning_rate = 0.0001\n",
    "    segment_num = None\n",
    "    best_train_path = None\n",
    "    best_val_path = None\n",
    "    vid_type = None\n",
    "    shifter = True\n",
    "    \n",
    "    \n",
    "args=Args()"
   ]
  },
  {
   "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",
    "               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",
    "               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=4)\n",
    "    val_dataloader = DataLoader(dataset=val_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)\n",
    "    \n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)\n",
    "\n",
    "    best_train_loss = np.inf\n",
    "    best_val_loss = np.inf\n",
    "    train_loss_list = []\n",
    "    val_loss_list = []\n",
    "\n",
    "    # start training\n",
    "    ct = 0\n",
    "    for epoch in range(args.epochs):\n",
    "\n",
    "        print(\"Start epoch\", epoch)\n",
    "\n",
    "        model.train()\n",
    "        print(sum([param.nelement() for param in model.parameters()]))\n",
    "\n",
    "\n",
    "        epoch_train_loss = 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",
    "            image = torch.squeeze(image, axis=1)\n",
    "\n",
    "            pred = model(image, behav)\n",
    "            \n",
    "            loss = nn.functional.poisson_nll_loss(pred, spikes, reduction='mean', log_input=False)\n",
    "\n",
    "            epoch_train_loss += loss.item()\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        epoch_train_loss = epoch_train_loss / len(train_dataloader)\n",
    "\n",
    "        train_loss_list.append(epoch_train_loss)\n",
    "\n",
    "        if epoch_train_loss < best_train_loss:\n",
    "\n",
    "            torch.save(model.state_dict(), args.best_train_path)\n",
    "            best_train_loss = epoch_train_loss\n",
    "\n",
    "        print(\"Epoch {} train loss: {}\".format(epoch, epoch_train_loss))\n",
    "\n",
    "        model.eval()\n",
    "\n",
    "        epoch_val_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",
    "                image = torch.squeeze(image, axis=1)\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_loss += loss.item()\n",
    "\n",
    "        epoch_val_loss = epoch_val_loss / len(val_dataloader)\n",
    "\n",
    "        val_loss_list.append(epoch_val_loss)\n",
    "\n",
    "        if epoch_val_loss < best_val_loss:\n",
    "            ct = 0\n",
    "            torch.save(model.state_dict(), args.best_val_path)\n",
    "            best_val_loss = epoch_val_loss\n",
    "        else: \n",
    "            ct+=1\n",
    "            if ct>=5: \n",
    "                print('stop trianing')\n",
    "                break\n",
    "\n",
    "        print(\"Epoch {} val loss: {}\".format(epoch, epoch_val_loss))\n",
    "\n",
    "        print(\"End epoch\", epoch)\n",
    "        \n",
    "    return train_loss_list, val_loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT vid_mean True\n",
      "30120 7540\n",
      "Start epoch 0\n",
      "4162958\n",
      "Epoch 0 train loss: 0.7149918049068774\n",
      "Epoch 0 val loss: 0.6612052838007609\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4162958\n",
      "Epoch 1 train loss: 0.628875284376791\n",
      "Epoch 1 val loss: 0.6400157968203227\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4162958\n",
      "Epoch 2 train loss: 0.613935808003959\n",
      "Epoch 2 val loss: 0.6331848402818044\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4162958\n",
      "Epoch 3 train loss: 0.6048753685870413\n",
      "Epoch 3 val loss: 0.6299521227677664\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4162958\n",
      "Epoch 4 train loss: 0.5973379601866512\n",
      "Epoch 4 val loss: 0.6300125598907471\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4162958\n",
      "Epoch 5 train loss: 0.5911980357210515\n",
      "Epoch 5 val loss: 0.6264310300350189\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4162958\n",
      "Epoch 6 train loss: 0.5852572791657206\n",
      "Epoch 6 val loss: 0.6293766975402832\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "4162958\n",
      "Epoch 7 train loss: 0.5786201297226599\n",
      "Epoch 7 val loss: 0.6309882203737894\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "4162958\n",
      "Epoch 8 train loss: 0.5724357300895756\n",
      "Epoch 8 val loss: 0.629713241259257\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "4162958\n",
      "Epoch 9 train loss: 0.5663872537976604\n",
      "Epoch 9 val loss: 0.632382728656133\n",
      "End epoch 9\n",
      "Start epoch 10\n",
      "4162958\n",
      "Epoch 10 train loss: 0.5607127511905412\n",
      "stop trianing\n",
      "070921_J553RT vid_mean False\n",
      "30120 7540\n",
      "Start epoch 0\n",
      "4162958\n",
      "Epoch 0 train loss: 0.7066458862716869\n",
      "Epoch 0 val loss: 0.650882093111674\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4162958\n",
      "Epoch 1 train loss: 0.6179687603045319\n",
      "Epoch 1 val loss: 0.6381854216257731\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4162958\n",
      "Epoch 2 train loss: 0.601754866919275\n",
      "Epoch 2 val loss: 0.6351901014645894\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4162958\n",
      "Epoch 3 train loss: 0.5891214932425547\n",
      "Epoch 3 val loss: 0.6368765274683634\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4162958\n",
      "Epoch 4 train loss: 0.5765726384470018\n",
      "Epoch 4 val loss: 0.6403344651063283\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4162958\n",
      "Epoch 5 train loss: 0.5632098903090267\n",
      "Epoch 5 val loss: 0.6463113327821096\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4162958\n",
      "Epoch 6 train loss: 0.5495615056005575\n",
      "Epoch 6 val loss: 0.6518809139728546\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "4162958\n",
      "Epoch 7 train loss: 0.5359701978958259\n",
      "stop trianing\n",
      "101521_J559NC vid_mean True\n",
      "42410 10610\n",
      "Start epoch 0\n",
      "4138619\n",
      "Epoch 0 train loss: 0.7661719397608057\n",
      "Epoch 0 val loss: 0.7241966823736826\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4138619\n",
      "Epoch 1 train loss: 0.7020898719868028\n",
      "Epoch 1 val loss: 0.7095749207905361\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4138619\n",
      "Epoch 2 train loss: 0.684949310429125\n",
      "Epoch 2 val loss: 0.7074893741380601\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4138619\n",
      "Epoch 3 train loss: 0.6741039214364017\n",
      "Epoch 3 val loss: 0.7049783085073743\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4138619\n",
      "Epoch 4 train loss: 0.6638500374483775\n",
      "Epoch 4 val loss: 0.7057554253510067\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4138619\n",
      "Epoch 5 train loss: 0.6542803767215775\n",
      "Epoch 5 val loss: 0.7056908295268104\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4138619\n",
      "Epoch 6 train loss: 0.6447855288005737\n",
      "Epoch 6 val loss: 0.706460687376204\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "4138619\n",
      "Epoch 7 train loss: 0.6352315153702196\n",
      "Epoch 7 val loss: 0.7129633937563215\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "4138619\n",
      "Epoch 8 train loss: 0.6265789306307413\n",
      "stop trianing\n",
      "101521_J559NC vid_mean False\n",
      "42410 10610\n",
      "Start epoch 0\n",
      "4138619\n",
      "Epoch 0 train loss: 0.7614599362195257\n",
      "Epoch 0 val loss: 0.7205266356468201\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4138619\n",
      "Epoch 1 train loss: 0.6877516500921135\n",
      "Epoch 1 val loss: 0.7138267897424244\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4138619\n",
      "Epoch 2 train loss: 0.6640911439815199\n",
      "Epoch 2 val loss: 0.7156213238125756\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4138619\n",
      "Epoch 3 train loss: 0.6418072057057576\n",
      "Epoch 3 val loss: 0.7247993875117529\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4138619\n",
      "Epoch 4 train loss: 0.6198560746319323\n",
      "Epoch 4 val loss: 0.7350810127598899\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4138619\n",
      "Epoch 5 train loss: 0.6002868543188256\n",
      "Epoch 5 val loss: 0.743152957587015\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4138619\n",
      "Epoch 6 train loss: 0.5836871550025711\n",
      "stop trianing\n",
      "110421_J569LT vid_mean True\n",
      "32940 8240\n",
      "Start epoch 0\n",
      "4116842\n",
      "Epoch 0 train loss: 0.7858602155086606\n",
      "Epoch 0 val loss: 0.7689404306989728\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4116842\n",
      "Epoch 1 train loss: 0.7262889807538468\n",
      "Epoch 1 val loss: 0.7570791930863352\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4116842\n",
      "Epoch 2 train loss: 0.7066463195076285\n",
      "Epoch 2 val loss: 0.7529976295702385\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4116842\n",
      "Epoch 3 train loss: 0.6929255751676338\n",
      "Epoch 3 val loss: 0.755905826886495\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4116842\n",
      "Epoch 4 train loss: 0.6830630953921828\n",
      "Epoch 4 val loss: 0.7541479396097588\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4116842\n",
      "Epoch 5 train loss: 0.6706245837285537\n",
      "Epoch 5 val loss: 0.7552830883950898\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4116842\n",
      "Epoch 6 train loss: 0.6621160336243089\n",
      "Epoch 6 val loss: 0.7556237361647866\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "4116842\n",
      "Epoch 7 train loss: 0.6512243114700613\n",
      "stop trianing\n",
      "110421_J569LT vid_mean False\n",
      "32940 8240\n",
      "Start epoch 0\n",
      "4116842\n",
      "Epoch 0 train loss: 0.7731393894483877\n",
      "Epoch 0 val loss: 0.7706828279928728\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "4116842\n",
      "Epoch 1 train loss: 0.7010577311811521\n",
      "Epoch 1 val loss: 0.7657317392753832\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "4116842\n",
      "Epoch 2 train loss: 0.674573653890181\n",
      "Epoch 2 val loss: 0.7668708237734708\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "4116842\n",
      "Epoch 3 train loss: 0.6507555893225263\n",
      "Epoch 3 val loss: 0.7769291581529559\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "4116842\n",
      "Epoch 4 train loss: 0.6279067272363708\n",
      "Epoch 4 val loss: 0.7869002782937252\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "4116842\n",
      "Epoch 5 train loss: 0.6064752497414286\n",
      "Epoch 5 val loss: 0.7993351582324866\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "4116842\n",
      "Epoch 6 train loss: 0.5873209399770397\n",
      "stop trianing\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# for file_id, num_neurons in [(\"110421_J569LT\", 52), (\"101521_J559NC\", 63), (\"070921_J553RT\", 108)]:\n",
    "for file_id, num_neurons in [(\"070921_J553RT\", 68), (\"101521_J559NC\", 49), (\"110421_J569LT\", 32)]:\n",
    "    for vid_type in [\"vid_mean\"]:\n",
    "        for shifter in [True, False]:\n",
    "            print(file_id, vid_type, shifter)\n",
    "            args.segment_num = 10\n",
    "            args.vid_type = vid_type\n",
    "            args.num_neurons = num_neurons\n",
    "            args.shifter = shifter\n",
    "            args.file_id = file_id\n",
    "            args.best_train_path = \"/hdd/yuchen/trainEffNetShifter{}_{}_{}_{}.pth\".format(args.shifter, args.segment_num, \n",
    "                                                                                            args.vid_type, \n",
    "                                                                                            args.file_id)\n",
    "            args.best_val_path = \"/hdd/yuchen/valEffNetShifter{}_{}_{}_{}.pth\".format(args.shifter, args.segment_num, \n",
    "                                                                                        args.vid_type, \n",
    "                                                                                        args.file_id)\n",
    "            \n",
    "            model = Predictor(num_neurons=args.num_neurons, k1=7, k2=7, k3=7).to(device)\n",
    "    \n",
    "            train_loss_list, val_loss_list = train_model()\n"
   ]
  },
  {
   "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=2)\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",
    "            image = torch.squeeze(image, axis=1)\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 vid_mean True\n",
      "30120 7540\n",
      "MSE 0.069431\n",
      "mean corr, 0.542+-0.153\n",
      "070921_J553RT vid_mean False\n",
      "30120 7540\n",
      "MSE 0.068346\n",
      "mean corr, 0.521+-0.153\n",
      "101521_J559NC vid_mean True\n",
      "42410 10610\n",
      "MSE 0.096501\n",
      "mean corr, 0.510+-0.127\n",
      "101521_J559NC vid_mean False\n",
      "42410 10610\n",
      "MSE 0.098369\n",
      "mean corr, 0.468+-0.145\n",
      "110421_J569LT vid_mean True\n",
      "32940 8240\n",
      "MSE 0.103219\n",
      "mean corr, 0.393+-0.165\n",
      "110421_J569LT vid_mean False\n",
      "32940 8240\n",
      "MSE 0.109311\n",
      "mean corr, 0.349+-0.183\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# for file_id, num_neurons in [(\"110421_J569LT\", 52), (\"101521_J559NC\", 63), (\"070921_J553RT\", 108)]:\n",
    "for file_id, num_neurons in [(\"070921_J553RT\", 68), (\"101521_J559NC\", 49), (\"110421_J569LT\", 32)]:\n",
    "\n",
    "    for vid_type in [\"vid_mean\"]:\n",
    "        for shifter in [True, False]:\n",
    "            args.shifter=shifter\n",
    "    \n",
    "            print(file_id, vid_type, shifter)\n",
    "    \n",
    "            args.segment_num = 10\n",
    "            args.vid_type = vid_type\n",
    "            args.num_neurons = num_neurons\n",
    "            args.file_id = file_id\n",
    "            args.best_train_path = \"/hdd/yuchen/trainEffNetShifter{}_{}_{}_{}.pth\".format(args.shifter, \n",
    "                                                                                           args.segment_num,\n",
    "                                                                                            args.vid_type, \n",
    "                                                                                            args.file_id)\n",
    "            args.best_val_path = \"/hdd/yuchen/valEffNetShifter{}_{}_{}_{}.pth\".format(args.shifter, \n",
    "                                                                                       args.segment_num, \n",
    "                                                                                        args.vid_type, \n",
    "                                                                                        args.file_id)\n",
    "            \n",
    "            model = Predictor(num_neurons=args.num_neurons, k1=7, k2=7, k3=7).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",
    "            \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)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}