saliency_map_cnn_gru_shifter_fig_5.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\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\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import Subset, DataLoader, ConcatDataset\n",
    "from mouse_model.data_utils_new import MouseDatasetSegNewBehav"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Shifter(nn.Module):\n",
    "    \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",
    "    \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\n",
    "\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=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": "markdown",
   "metadata": {},
   "source": [
    "set args & random seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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",
    "def set_random_seed(seed: int, deterministic: bool = True):\n",
    "    # from nnfabrik package\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    if deterministic:\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "    torch.manual_seed(seed)  # this sets both CPU and CUDA seeds for PyTorch\n",
    "\n",
    "seed = args.seed\n",
    "set_random_seed(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",
    "\n",
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "saliency map function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "device=\"cuda\"\n",
    "\n",
    "def get_saliency_map(model, dataloader, neuron_idx, device=\"cuda\"):\n",
    "    \n",
    "    grad_list = []\n",
    "    \n",
    "    for (image, behav, spikes) in dataloader:\n",
    "        \n",
    "        \n",
    "        image, behav, spikes = image.to(device), behav.to(device), spikes.to(device)\n",
    "            \n",
    "        x = Variable(behav, requires_grad=True)\n",
    "\n",
    "        output = model(image, x)\n",
    "        output = torch.sum(output[:, neuron_idx]) # can use different batch sizes in the dataloader\n",
    "        output.backward()\n",
    "        \n",
    "        grad_list.append(x.grad.to('cpu').numpy())\n",
    "        \n",
    "    grad_np = np.concatenate(grad_list, axis=0)\n",
    "\n",
    "    weights = np.mean(np.squeeze(grad_np), axis=0)\n",
    "        \n",
    "    return weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": "markdown",
   "metadata": {},
   "source": [
    "Mouse 1: 070921_J553RT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.file_id = \"070921_J553RT\"\n",
    "args.num_neurons = 68\n",
    "args.shifter = True\n",
    "\n",
    "args.behav_mode = \"all_prod\"\n",
    "args.behav_dim = 66\n",
    "\n",
    "args.seq_len = 1\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",
    "weights_path = \"weights_cnn_gru_shifter/val_070921_J553RT_all_prod_seq_1.pth\"\n",
    "\n",
    "model.load_state_dict(torch.load(weights_path))\n",
    "# freeze layer param\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "# freeze batchnorm statistics\n",
    "model.shifter.layers[0].eval()\n",
    "model.shifter.layers[2].eval()\n",
    "model.shifter.layers[5].eval()\n",
    "model.visual_encoder.layers[1].eval()\n",
    "model.visual_encoder.layers[5].eval()\n",
    "model.visual_encoder.layers[9].eval()\n",
    "model.behav_encoder.layers[0].eval()\n",
    "model.bn.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = load_test_ds()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataloader = DataLoader(dataset=test_ds, batch_size=256, shuffle=False, num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "for neuron_idx in range(68):\n",
    "    res = get_saliency_map(model, test_dataloader, neuron_idx=neuron_idx, device=device)\n",
    "    np.save(\"saliency_map_cnn_gru_shifter/070921_J553RT/{}.npy\".format(neuron_idx), res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mouse 3: 101521_J559NC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.file_id = \"101521_J559NC\"\n",
    "args.num_neurons = 49\n",
    "args.shifter = True\n",
    "\n",
    "args.behav_mode = \"all_prod\"\n",
    "args.behav_dim = 66\n",
    "\n",
    "args.seq_len = 1\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",
    "weights_path = \"weights_cnn_gru_shifter/val_101521_J559NC_all_prod_seq_1.pth\"\n",
    "\n",
    "model.load_state_dict(torch.load(weights_path))\n",
    "# freeze layer param\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "# freeze batchnorm statistics\n",
    "model.shifter.layers[0].eval()\n",
    "model.shifter.layers[2].eval()\n",
    "model.shifter.layers[5].eval()\n",
    "model.visual_encoder.layers[1].eval()\n",
    "model.visual_encoder.layers[5].eval()\n",
    "model.visual_encoder.layers[9].eval()\n",
    "model.behav_encoder.layers[0].eval()\n",
    "model.bn.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = load_test_ds()\n",
    "test_dataloader = DataLoader(dataset=test_ds, batch_size=256, shuffle=False, num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "for neuron_idx in range(49):\n",
    "    res = get_saliency_map(model, test_dataloader, neuron_idx=neuron_idx, device=device)\n",
    "    np.save(\"saliency_map_cnn_gru_shifter/101521_J559NC/{}.npy\".format(neuron_idx), res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mouse 2: 110421_J569LT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.file_id = \"110421_J569LT\"\n",
    "args.num_neurons = 32\n",
    "args.shifter = True\n",
    "\n",
    "args.behav_mode = \"all_prod\"\n",
    "args.behav_dim = 66\n",
    "\n",
    "args.seq_len = 1\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",
    "weights_path = \"weights_cnn_gru_shifter/val_110421_J569LT_all_prod_seq_1.pth\"\n",
    "\n",
    "model.load_state_dict(torch.load(weights_path))\n",
    "# freeze layer param\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "# freeze batchnorm statistics\n",
    "model.shifter.layers[0].eval()\n",
    "model.shifter.layers[2].eval()\n",
    "model.shifter.layers[5].eval()\n",
    "model.visual_encoder.layers[1].eval()\n",
    "model.visual_encoder.layers[5].eval()\n",
    "model.visual_encoder.layers[9].eval()\n",
    "model.behav_encoder.layers[0].eval()\n",
    "model.bn.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = load_test_ds()\n",
    "test_dataloader = DataLoader(dataset=test_ds, batch_size=256, shuffle=False, num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "for neuron_idx in range(32):\n",
    "    res = get_saliency_map(model, test_dataloader, neuron_idx=neuron_idx, device=device)\n",
    "    np.save(\"saliency_map_cnn_gru_shifter/110421_J569LT/{}.npy\".format(neuron_idx), res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}