bench_vsr_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",
    "epoch_acc = {}\n",
    "\n",
    "dataloader = zip(*test_loader.values())\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch_index, multitask_batch in tqdm(enumerate(dataloader)):\n",
    "        stim_batch, resp_batch, seq_len, list_length = multitask_batch[6]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        for model_name, model in model_dict.items():\n",
    "            out, _, _, _, _ = model(stim_batch, 'VSR_Task', seq_len)\n",
    "            pred = torch.argmax(out, dim=-1)\n",
    "\n",
    "            for index, _ in enumerate(seq_len):\n",
    "                curr_gt = resp_batch[index, list_length[index]:list_length[index]*2]\n",
    "                curr_pred = pred[index, list_length[index]:list_length[index]*2]\n",
    "\n",
    "                for count, item in enumerate(curr_gt):\n",
    "                    if (model_name+'_List_Length_'+str(list_length[index].item())+\n",
    "                        '_SP_'+str(count+1)) not in epoch_acc:\n",
    "                        epoch_acc[model_name+'_List_Length_'+\n",
    "                                  str(list_length[index].item())+'_SP_'+str(count+1)] = []\n",
    "\n",
    "                    if item == curr_pred[count]:\n",
    "                        epoch_acc[model_name+'_List_Length_'+\n",
    "                                  str(list_length[index].item())+'_SP_'+str(count+1)].append(1)\n",
    "                    else:\n",
    "                        epoch_acc[model_name+'_List_Length_'+\n",
    "                                  str(list_length[index].item())+'_SP_'+str(count+1)].append(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('vsr_accs.json', 'w') as fp:\n",
    "    json.dump(epoch_acc, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import stats\n",
    "\n",
    "epoch_acc_std_err = {}\n",
    "\n",
    "for key in epoch_acc:\n",
    "    epoch_acc_std_err[key] = [np.mean(epoch_acc[key]), \n",
    "                              stats.sem(epoch_acc[key])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = {}\n",
    "\n",
    "for key in list(epoch_acc_std_err.keys()):\n",
    "    model = key.split('_')[0] + '_' + key.split('_')[1]\n",
    "    list_length = key.split('_')[4]\n",
    "    sp = key.split('_')[6]\n",
    "\n",
    "    if model not in plot_data:\n",
    "        plot_data[model] = {}\n",
    "    \n",
    "    if list_length not in plot_data[model]:\n",
    "        plot_data[model][list_length] = [[], [], []]\n",
    "\n",
    "    plot_data[model][list_length][0].append(sp)\n",
    "    plot_data[model][list_length][1].append(epoch_acc_std_err[key][0])\n",
    "    plot_data[model][list_length][2].append(epoch_acc_std_err[key][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "human_plot_data = {\n",
    "    '3': [['1', '2', '3'], [7.16/8, 6.8/8, 6.9/8], []], \n",
    "    '4': [['1', '2', '3', '4'], [6.35/8, 5.25/8, 4.8/8, 6.05/8], []],\n",
    "    '5': [['1', '2', '3', '4', '5'], [5.5/8, 4.3/8, 4.2/8, 3.6/8, 5/8], []], \n",
    "    '6': [['1', '2', '3', '4', '5', '6'], [5.1/8, 4.25/8, 3.55/8, 3.5/8, 2.9/8, 3.95/8], []]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "from matplotlib import cm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Serif\"\n",
    "sns.set_context(\"paper\", font_scale=1.5, rc={\"lines.linewidth\": 1})\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "# Create twp subplots\n",
    "ax1 = plt.subplot(2, 4, 7)\n",
    "ax2 = plt.subplot(1, 2, 1)\n",
    "\n",
    "\n",
    "for model in list(plot_data.keys()):\n",
    "    if model == 'gru_1024':\n",
    "        for list_length in list(plot_data[model].keys()):\n",
    "            if list_length in ['3', '4', '5', '6', '7', '9']:\n",
    "                ax2.errorbar(plot_data[model][list_length][0], \n",
    "                            plot_data[model][list_length][1], \n",
    "                            yerr=plot_data[model][list_length][2],\n",
    "                            label='List Length '+list_length, fmt='o-',\n",
    "                            linewidth=2, markersize=3, capsize=4)\n",
    "                \n",
    "                \n",
    "ax2.set_ylabel('Top-1 Accuracy', fontsize=25)\n",
    "\n",
    "ax2.set_xticks([0, 1, 2, 3, 4, 5, 6, 7, 8])\n",
    "ax2.set_xticklabels([1, 2, 3, 4, 5, 6, 7, 8, 9])\n",
    "\n",
    "ax2.set_ylim([0.5, 1.00])\n",
    "ax2.set_yticks([0.5, 0.6, 0.7, 0.8, 0.9, 1.0])\n",
    "\n",
    "\n",
    "for list_length in list(human_plot_data.keys()):        \n",
    "    ax1.errorbar(human_plot_data[list_length][0], \n",
    "                human_plot_data[list_length][1], \n",
    "                yerr=0, \n",
    "                label='List Length '+list_length, fmt='o-',\n",
    "                linewidth=2, markersize=7, color='C'+str(int(list_length)-2))\n",
    "            \n",
    "sns.despine(left=False, bottom=False, right=True, top=True)\n",
    "\n",
    "\n",
    "ax1.set_xticks([0, 1, 2, 3, 4, 5])\n",
    "ax1.set_xticklabels([1, 2, 3, 4, 5, 6], fontsize=18)\n",
    "\n",
    "ax1.set_yticks([0.4, 0.6, 0.8, 1.0])\n",
    "ax1.set_yticklabels([0.4, 0.6, 0.8, 1.0], fontsize=18)\n",
    "\n",
    "ax1.set_ylim([0.3, 1.01])\n",
    "\n",
    "\n",
    "plt.xticks(fontsize=20)\n",
    "plt.yticks(fontsize=20)\n",
    "\n",
    "plt.legend(frameon=False, loc='upper center', \n",
    "           bbox_to_anchor=(1.33, 1.05), ncol=1, prop={'size': 16})\n",
    "\n",
    "plt.text(5, 0.405, 'Serial Position', fontsize=25)\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "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
}