bench_sts_task.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "from argparse import Namespace\n",
    "\n",
    "\n",
    "config = Namespace(\n",
    "    data_folder='./wm_bench_data', \n",
    "    max_seq_len=20, \n",
    "    rs_img_size=32, \n",
    "    batch_size=10, \n",
    "    num_workers=4, \n",
    "    use_cnn=1, \n",
    "    model_path='./model.pt'\n",
    ")\n",
    "\n",
    "device = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.utils.data_utils import get_test_multitask_dataloader\n",
    "\n",
    "test_loader = get_test_multitask_dataloader(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.model import WM_Model\n",
    "\n",
    "model_data = torch.load(config.model_path)\n",
    "model = WM_Model(Namespace(**model_data['config']), device).to(device)\n",
    "model.load_state_dict(model_data['model_state_dict'])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = {\n",
    "    'model_1': model,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rnn_out_all = {}\n",
    "\n",
    "resp_batch_all = []\n",
    "\n",
    "dataloader = zip(*test_loader.values())\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch_index, multi_task_batch in tqdm(enumerate(dataloader)):\n",
    "        stim_batch, resp_batch, seq_len = multi_task_batch[4]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        resp_batch_all.append(resp_batch.cpu().numpy())\n",
    "\n",
    "        for model_name, model in model_dict.items():\n",
    "            out, rnn_out, hn, proj_out, _ = model(stim_batch, 'STSC_Task', seq_len)\n",
    "            rnn_out = rnn_out.cpu().numpy()\n",
    "\n",
    "            if model_name not in rnn_out_all:\n",
    "                rnn_out_all[model_name] = [rnn_out]\n",
    "            else:\n",
    "                rnn_out_all[model_name].append(rnn_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp_batch_all = np.concatenate(resp_batch_all, axis=0)\n",
    "resp_batch_all = resp_batch_all.reshape(-1)\n",
    "\n",
    "for model_name, rnn_out in rnn_out_all.items():\n",
    "    rnn_out_all[model_name] = np.concatenate(rnn_out, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "test_data = json.load(open('./wm_bench_data/Spatial_Task_Switching_Cued/data_rand_trials.json', \n",
    "                           'r'))\n",
    "\n",
    "tasks_gt_map = {'Up_Down': 0, 'Left_Right': 1, 'Cue_Up_Down': 2, 'Cue_Left_Right': 3}\n",
    "\n",
    "tasks_gt = []\n",
    "for trial in test_data[\"test\"]:\n",
    "    for task in trial['task_gt']:\n",
    "        tasks_gt.append(tasks_gt_map[task])\n",
    "\n",
    "tasks_gt = np.array(tasks_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_states = {}\n",
    "\n",
    "hidden_states['lstm_1024'] = rnn_out_all['lstm_1024'].reshape(-1, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasksss_gt = []\n",
    "respsss_gt = []\n",
    "hidden_statesss = []\n",
    "\n",
    "for index, val in enumerate(resp_batch_all):\n",
    "    if val != 2:\n",
    "        tasksss_gt.append(tasks_gt[index])\n",
    "        respsss_gt.append(val)\n",
    "        hidden_statesss.append(hidden_states['lstm_1024'][index])\n",
    "\n",
    "tasksss_gt = np.array(tasksss_gt)\n",
    "respsss_gt = np.array(respsss_gt)\n",
    "hidden_statesss = np.array(hidden_statesss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "pca = PCA(n_components=2, random_state=69)\n",
    "pca.fit(hidden_statesss)\n",
    "\n",
    "hidden_statesss_pca = pca.transform(hidden_statesss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "\n",
    "sns.set_style(\"ticks\")\n",
    "\n",
    "# Change font style\n",
    "plt.rcParams[\"font.family\"] = \"Serif\"\n",
    "sns.set_context(\"paper\", font_scale=1.5, rc={\"lines.linewidth\": 1})\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "\n",
    "scatterplot = sns.scatterplot(x=hidden_statesss_pca[:, 0],\n",
    "                y=hidden_statesss_pca[:, 1], \n",
    "                hue=respsss_gt, hue_order=[0, 1], edgecolor='none', size=0.2, palette='tab10')\n",
    "\n",
    "sns.despine(left=False, bottom=False, right=True, top=True)\n",
    "\n",
    "handles, labels = scatterplot.get_legend_handles_labels()\n",
    "\n",
    "# Define custom labels for the legend\n",
    "legend_labels = ['Left / Top GT', 'Right / Bottom GT']\n",
    "\n",
    "# Create a new legend with custom labels\n",
    "scatterplot.legend(handles=handles, frameon=False, bbox_to_anchor=(0.09, 0.98), \n",
    "                   labels=legend_labels, fontsize=20, markerscale=2)\n",
    "\n",
    "\n",
    "plt.xlabel('PC1', fontsize=25)\n",
    "plt.ylabel('PC2', fontsize=25)\n",
    "\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "# Plot the hidden states in scatter plot using seaborn\n",
    "\n",
    "sns.set_style(\"ticks\")\n",
    "\n",
    "# Change font style\n",
    "plt.rcParams[\"font.family\"] = \"Serif\"\n",
    "sns.set_context(\"paper\", font_scale=1.5, rc={\"lines.linewidth\": 1})\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "\n",
    "scatterplot = sns.scatterplot(x=hidden_statesss_pca[:, 0],\n",
    "                y=hidden_statesss_pca[:, 1], \n",
    "                hue=tasksss_gt, hue_order=[0, 1], edgecolor='none', size=0.2, \n",
    "                palette=['tab:green','tab:red'])\n",
    "\n",
    "\n",
    "sns.despine(left=False, bottom=False, right=True, top=True)\n",
    "\n",
    "handles, labels = scatterplot.get_legend_handles_labels()\n",
    "\n",
    "# Define custom labels for the legend\n",
    "legend_labels = ['Top vs Bottom Task', 'Left vs Right Task']\n",
    "\n",
    "# Create a new legend with custom labels\n",
    "scatterplot.legend(handles=handles, frameon=False, bbox_to_anchor=(0.09, 0.98), \n",
    "                   labels=legend_labels, fontsize=20, markerscale=2)\n",
    "\n",
    "\n",
    "plt.xlabel('PC1', fontsize=25)\n",
    "plt.ylabel('PC2', fontsize=25)\n",
    "\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memnet",
   "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.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}