train_autoencoder_shifter_table_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\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, os\n",
    "from kornia.geometry.transform import get_affine_matrix2d, warp_affine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the model"
   ]
  },
  {
   "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, c1=32, c2=64, c3=128, dropout_prob=0):\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=64, kernel_size=k1, stride=2),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=k2, stride=2),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=k3, stride=2),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(self.output_shape[0]*self.output_shape[1]*c3, output_dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.layers(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VisualDecoder777(nn.Module):\n",
    "    \n",
    "    # input_dim is latent size\n",
    "    def __init__(self, input_dim, c1=128, c2=64, c3=32, dropout_prob=0):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.input_dim = input_dim\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Linear(input_dim, c1*3*5),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.Unflatten(1, (c1, 3, 5)),\n",
    "#             PrintLayer(),\n",
    "            nn.ConvTranspose2d(in_channels=c1, out_channels=c2, kernel_size=7, stride=2, output_padding=(0,1)),\n",
    "#             PrintLayer(),\n",
    "            nn.BatchNorm2d(c2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.ConvTranspose2d(in_channels=c2, out_channels=c3, kernel_size=7, stride=2),\n",
    "#             PrintLayer(),\n",
    "            nn.BatchNorm2d(c3),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.ConvTranspose2d(in_channels=c3, out_channels=1, kernel_size=7, stride=2, output_padding=1),\n",
    "#             PrintLayer(),\n",
    "            nn.BatchNorm2d(1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "\n",
    "        x = self.conv_layers(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Shifter(nn.Module):\n",
    "    def __init__(self, input_dim=4, output_dim=3, hidden_dim=128):\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\n",
    "        \n",
    "class AutoencoderPredictor777(nn.Module):\n",
    "    \n",
    "    def __init__(self, num_neurons, latent_size, c1=32, c2=64, c3=128, dropout_prob=0):\n",
    "\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = VisualEncoder(output_dim=latent_size, k1=7, k2=7, k3=7, \n",
    "                                     c1=c1, c2=c2, c3=c3, dropout_prob=dropout_prob)\n",
    "        self.decoder = VisualDecoder777(input_dim=latent_size, c1=c3, c2=c2, c3=c1, dropout_prob=dropout_prob)\n",
    "        self.fc = nn.Linear(latent_size, num_neurons)\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(\n",
    "                                            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",
    "        latent_vec = self.encoder(images)\n",
    "        recon = self.decoder(latent_vec)\n",
    "        pred = self.fc(latent_vec)\n",
    "        pred = self.softplus(pred)\n",
    "        \n",
    "        return pred, recon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "class Args:\n",
    "    \n",
    "    seed = 0\n",
    "    batch_size = 256\n",
    "    learning_rate = 0.0001\n",
    "    epochs = 50\n",
    "    alpha = 0.5\n",
    "    file_id = None\n",
    "    num_neurons = None\n",
    "    behav_mode = \"orig\"\n",
    "    best_val_path = None\n",
    "    best_train_path = None\n",
    "    vid_type = \"vid_mean\"\n",
    "    segment_num = 10\n",
    "    seq_len = 1\n",
    "    vis_latent_dim = 256\n",
    "    shifter=True\n",
    "    \n",
    "args=Args()\n",
    "\n",
    "seed = args.seed\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\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",
    "\n",
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_train_val_ds():\n",
    "    ds_list = [MouseDatasetSegNewBehav(\n",
    "        file_id=args.file_id, seg_idx=i, segment_num=args.segment_num, data_split=\"train\", vid_type=args.vid_type, \n",
    "        seq_len=args.seq_len, predict_offset=1, 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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_test_ds():\n",
    "    test_ds = [MouseDatasetSegNewBehav(\n",
    "        file_id=args.file_id, seg_idx=i, segment_num=args.segment_num, data_split=\"test\", vid_type=args.vid_type, \n",
    "        seq_len=args.seq_len, predict_offset=1, 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": 8,
   "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",
    "    ct = 0\n",
    "    \n",
    "    # start training\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, epoch_train_recon_loss = 0, 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",
    "            image = torch.squeeze(image, axis=1)\n",
    "            \n",
    "            pred, recon = model(image, behav)\n",
    "            \n",
    "            recon_loss = torch.nn.functional.mse_loss(recon, image, reduction='mean')\n",
    "            spike_loss = nn.functional.poisson_nll_loss(pred, spikes, reduction='mean', log_input=False)\n",
    "            total_loss = spike_loss + args.alpha * recon_loss\n",
    "\n",
    "            epoch_train_loss += total_loss.item()\n",
    "            epoch_train_spike_loss += spike_loss.item()\n",
    "            epoch_train_recon_loss += recon_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",
    "        epoch_train_recon_loss = epoch_train_recon_loss / len(train_dataloader)\n",
    "\n",
    "        train_loss_list.append([epoch_train_loss, epoch_train_spike_loss, epoch_train_recon_loss])\n",
    "        \n",
    "        print(\"Epoch {} train loss: {}, spike loss: {}, recon loss: {}\".format(\n",
    "            epoch, epoch_train_loss, epoch_train_spike_loss, epoch_train_recon_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_loss, epoch_val_spike_loss, epoch_val_recon_loss = 0, 0, 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, recon = model(image, behav)\n",
    "                \n",
    "                recon_loss = torch.nn.functional.mse_loss(recon, image, reduction='mean')\n",
    "                spike_loss = nn.functional.poisson_nll_loss(pred, spikes, reduction='mean', log_input=False)\n",
    "                total_loss = spike_loss + args.alpha * recon_loss\n",
    "\n",
    "                epoch_val_loss += total_loss.item()\n",
    "                epoch_val_spike_loss += spike_loss.item()\n",
    "                epoch_val_recon_loss += recon_loss.item()\n",
    "                \n",
    "        epoch_val_loss = epoch_val_loss / len(val_dataloader)\n",
    "        epoch_val_spike_loss = epoch_val_spike_loss / len(val_dataloader)\n",
    "        epoch_val_recon_loss = epoch_val_recon_loss / len(val_dataloader)\n",
    "\n",
    "        val_loss_list.append([epoch_val_loss, epoch_val_spike_loss, epoch_val_recon_loss])\n",
    "        \n",
    "        print(\"Epoch {} val spike loss: {}\".format(epoch, epoch_val_spike_loss))\n",
    "\n",
    "        if epoch_val_spike_loss < best_val_spike_loss:\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",
    "            ct = 0\n",
    "            \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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT\n",
      "30120 7540\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 0.6891603717359446, spike loss: 0.6473932842076835, recon loss: 0.08353417758214271\n",
      "save train model at epoch 0\n",
      "Epoch 0 val spike loss: 0.6400277396043141\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.6432702960604328, spike loss: 0.6129921649472189, recon loss: 0.06055626538345369\n",
      "save train model at epoch 1\n",
      "Epoch 1 val spike loss: 0.6308337152004242\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.6315709898027323, spike loss: 0.6025065148280839, recon loss: 0.0581289510858261\n",
      "save train model at epoch 2\n",
      "Epoch 2 val spike loss: 0.6292569975058238\n",
      "save val model at epoch 2\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.6232444840972706, spike loss: 0.5947218421152083, recon loss: 0.057045281217512435\n",
      "save train model at epoch 3\n",
      "Epoch 3 val spike loss: 0.6310473521550496\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.6158913025411509, spike loss: 0.5878888681783514, recon loss: 0.05600486954642554\n",
      "save train model at epoch 4\n",
      "Epoch 4 val spike loss: 0.6271923820177714\n",
      "save val model at epoch 4\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.6088746760861349, spike loss: 0.5812840002068018, recon loss: 0.055181349264615674\n",
      "save train model at epoch 5\n",
      "Epoch 5 val spike loss: 0.6312132835388183\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.6016679173808986, spike loss: 0.5744724687883409, recon loss: 0.054390900941976045\n",
      "save train model at epoch 6\n",
      "Epoch 6 val spike loss: 0.631647930542628\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.594711492627354, spike loss: 0.56774025195736, recon loss: 0.053942482160814736\n",
      "save train model at epoch 7\n",
      "Epoch 7 val spike loss: 0.6313550273577372\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "Epoch 8 train loss: 0.5882987314361637, spike loss: 0.561477168636807, recon loss: 0.05364312673524275\n",
      "save train model at epoch 8\n",
      "Epoch 8 val spike loss: 0.6378750681877137\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "Epoch 9 train loss: 0.5814901711577076, spike loss: 0.5547877277358103, recon loss: 0.05340488693850525\n",
      "save train model at epoch 9\n",
      "Epoch 9 val spike loss: 0.635291850566864\n",
      "End epoch 9\n",
      "Start epoch 10\n",
      "Epoch 10 train loss: 0.574539495221639, spike loss: 0.5479324298389887, recon loss: 0.05321413306992943\n",
      "save train model at epoch 10\n",
      "Epoch 10 val spike loss: 0.6480502625306447\n",
      "stop training\n",
      "110421_J569LT\n",
      "32940 8240\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 0.7862435930459074, spike loss: 0.7349134706711584, recon loss: 0.10266023989795714\n",
      "save train model at epoch 0\n",
      "Epoch 0 val spike loss: 0.751306521170067\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.7427434967469799, spike loss: 0.7030580593634022, recon loss: 0.07937088169792826\n",
      "save train model at epoch 1\n",
      "Epoch 1 val spike loss: 0.7471338058962966\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.7259001856626466, spike loss: 0.6885471399440322, recon loss: 0.07470608964677929\n",
      "save train model at epoch 2\n",
      "Epoch 2 val spike loss: 0.7479705575740698\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.7126357749451039, spike loss: 0.6767884880997413, recon loss: 0.0716945723045704\n",
      "save train model at epoch 3\n",
      "Epoch 3 val spike loss: 0.749487067713882\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.7016615031301513, spike loss: 0.6666219081989554, recon loss: 0.07007919072873833\n",
      "save train model at epoch 4\n",
      "Epoch 4 val spike loss: 0.7572717829184099\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.6924791451572447, spike loss: 0.6579118879266488, recon loss: 0.06913451492324356\n",
      "save train model at epoch 5\n",
      "Epoch 5 val spike loss: 0.7582550030766111\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.6826599928759789, spike loss: 0.6484224699264349, recon loss: 0.06847503948812336\n",
      "save train model at epoch 6\n",
      "Epoch 6 val spike loss: 0.7585058880574775\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.6736283755117609, spike loss: 0.6397129232569259, recon loss: 0.06783089838748754\n",
      "save train model at epoch 7\n",
      "Epoch 7 val spike loss: 0.7595067042292971\n",
      "stop training\n",
      "101521_J559NC\n",
      "42410 10610\n",
      "Start epoch 0\n",
      "Epoch 0 train loss: 0.7604727371629462, spike loss: 0.714831622609173, recon loss: 0.09128223260842173\n",
      "save train model at epoch 0\n",
      "Epoch 0 val spike loss: 0.7064318287940252\n",
      "save val model at epoch 0\n",
      "End epoch 0\n",
      "Start epoch 1\n",
      "Epoch 1 train loss: 0.7188562854944941, spike loss: 0.6834917477814548, recon loss: 0.07072907439376934\n",
      "save train model at epoch 1\n",
      "Epoch 1 val spike loss: 0.6987929088728768\n",
      "save val model at epoch 1\n",
      "End epoch 1\n",
      "Start epoch 2\n",
      "Epoch 2 train loss: 0.7038623787552477, spike loss: 0.6703905025160456, recon loss: 0.06694375436349087\n",
      "save train model at epoch 2\n",
      "Epoch 2 val spike loss: 0.6997889833790916\n",
      "End epoch 2\n",
      "Start epoch 3\n",
      "Epoch 3 train loss: 0.6925246374434736, spike loss: 0.6598558716745262, recon loss: 0.06533753467970584\n",
      "save train model at epoch 3\n",
      "Epoch 3 val spike loss: 0.6979122530846369\n",
      "save val model at epoch 3\n",
      "End epoch 3\n",
      "Start epoch 4\n",
      "Epoch 4 train loss: 0.682728387146111, spike loss: 0.6505466408758278, recon loss: 0.06436349379729076\n",
      "save train model at epoch 4\n",
      "Epoch 4 val spike loss: 0.698932735692887\n",
      "End epoch 4\n",
      "Start epoch 5\n",
      "Epoch 5 train loss: 0.6745941832841161, spike loss: 0.6427015680146505, recon loss: 0.06378522890070117\n",
      "save train model at epoch 5\n",
      "Epoch 5 val spike loss: 0.7012777371065957\n",
      "End epoch 5\n",
      "Start epoch 6\n",
      "Epoch 6 train loss: 0.6653881633138082, spike loss: 0.6337171265159744, recon loss: 0.06334207301218826\n",
      "save train model at epoch 6\n",
      "Epoch 6 val spike loss: 0.7017359194301424\n",
      "End epoch 6\n",
      "Start epoch 7\n",
      "Epoch 7 train loss: 0.6584246991628624, spike loss: 0.6269485243113644, recon loss: 0.062952354168856\n",
      "save train model at epoch 7\n",
      "Epoch 7 val spike loss: 0.704973544393267\n",
      "End epoch 7\n",
      "Start epoch 8\n",
      "Epoch 8 train loss: 0.649806940771011, spike loss: 0.6184882471360356, recon loss: 0.06263738453208682\n",
      "save train model at epoch 8\n",
      "Epoch 8 val spike loss: 0.70534234387534\n",
      "End epoch 8\n",
      "Start epoch 9\n",
      "Epoch 9 train loss: 0.641090191631432, spike loss: 0.6098934670528734, recon loss: 0.06239344136991415\n",
      "save train model at epoch 9\n",
      "Epoch 9 val spike loss: 0.708274013939358\n",
      "stop training\n"
     ]
    }
   ],
   "source": [
    "# 64 128 256\n",
    "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 shifter in [True]:\n",
    "        print(file_id)\n",
    "        args.vid_type = 'vid_mean'\n",
    "        args.file_id = file_id\n",
    "        args.num_neurons = num_neurons\n",
    "        args.shifter=shifter\n",
    "        args.best_train_path = \"/hdd/yuchen/trainAEshifter{}_{}_{}_{}.pth\".format(args.shifter, 1, args.vid_type, file_id)\n",
    "        args.best_val_path = \"/hdd/yuchen/valAEshifter{}_{}_{}_{}.pth\".format(args.shifter, 1, args.vid_type, file_id)\n",
    "    \n",
    "        model = AutoencoderPredictor777(num_neurons=args.num_neurons, \n",
    "                                        latent_size=args.vis_latent_dim,\n",
    "                                        c1=64, c2=128, c3=256).to(device)\n",
    "\n",
    "        train_loss_list, val_loss_list = train_model()\n",
    "\n",
    "    # np.save(\"/hdd/aiwenxu/net_weights_final/vis_autoencoder/loss_777_c1_64_{}.npy\".format(file_id), \n",
    "    #         np.array([train_loss_list, val_loss_list]))"
   ]
  },
  {
   "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",
    "            image = torch.squeeze(image, axis=1)\n",
    "            \n",
    "            pred, recon = 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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "070921_J553RT\n",
      "30120 7540\n",
      "R2 0.777193\n",
      "MSE 0.071034\n",
      "mean corr, 0.555+-0.140\n",
      "max corr 0.816352\n",
      "min corr 0.177165\n",
      "070921_J553RT\n",
      "30120 7540\n",
      "R2 0.765927\n",
      "MSE 0.076842\n",
      "mean corr, 0.509+-0.144\n",
      "max corr 0.773564\n",
      "min corr 0.115169\n",
      "110421_J569LT\n",
      "32940 8240\n",
      "R2 0.224810\n",
      "MSE 0.112353\n",
      "mean corr, 0.370+-0.145\n",
      "max corr 0.651283\n",
      "min corr 0.048037\n",
      "110421_J569LT\n",
      "32940 8240\n",
      "R2 0.210233\n",
      "MSE 0.114998\n",
      "mean corr, 0.375+-0.148\n",
      "max corr 0.661694\n",
      "min corr 0.110536\n",
      "101521_J559NC\n",
      "42410 10610\n",
      "R2 0.727099\n",
      "MSE 0.097395\n",
      "mean corr, 0.521+-0.144\n",
      "max corr 0.822660\n",
      "min corr 0.242436\n",
      "101521_J559NC\n",
      "42410 10610\n",
      "R2 0.663554\n",
      "MSE 0.117084\n",
      "mean corr, 0.459+-0.140\n",
      "max corr 0.749090\n",
      "min corr 0.166779\n"
     ]
    }
   ],
   "source": [
    "# 64 128 256\n",
    "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",
    "    \n",
    "    for shifter in [True, False]:\n",
    "        print(file_id)\n",
    "        args.vid_type = 'vid_mean'\n",
    "        args.file_id = file_id\n",
    "        args.num_neurons = num_neurons\n",
    "        args.shifter=shifter\n",
    "        args.best_train_path = \"/hdd/yuchen/trainAEshifter{}_{}_{}_{}.pth\".format(args.shifter, 1, args.vid_type, file_id)\n",
    "        args.best_val_path = \"/hdd/yuchen/valAEshifter{}_{}_{}_{}.pth\".format(args.shifter, 1, args.vid_type, file_id)\n",
    "        model = AutoencoderPredictor777(num_neurons=args.num_neurons, \n",
    "                                    latent_size=args.vis_latent_dim,\n",
    "                                    c1=64, c2=128, c3=256).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)))"
   ]
  },
  {
   "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
}