notebooks/pear-training-experiment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## PEAR Training Experiment\n",
    "\n",
    "This is a notebook where we train two models with different PEAR loss hyperparamters and compare their performance and consistency metrics.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from torch.optim import AdamW\n",
    "\n",
    "import pear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Define constants\n",
    "dataset = \"californiahousing\"\n",
    "batch_size = 64\n",
    "model_cfg = {\"name\": \"mlp\",\n",
    "             \"width\": 100,\n",
    "             \"depth\": 3}\n",
    "lr = 5e-4\n",
    "weight_decay = 2e-4\n",
    "explainers = [\"vanilla_gradients\", \"integrated_gradients\"]\n",
    "disagreement_mu = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "californiahousing dataset with 15475 training samples and 5159 testing samples and 8 features and (tensor([0, 1]), tensor([7686, 7789])) classes.\n"
     ]
    }
   ],
   "source": [
    "loader_train, loader_test, num_classes = pear.get_data(dataset,\n",
    "                                                       batch_size,\n",
    "                                                       data_path=\"../pear/datasets\")\n",
    "input_dim = loader_train.dataset.data.shape[1]\n",
    "num_training_data = loader_train.dataset.data.shape[0]\n",
    "num_testing_data = loader_test.dataset.data.shape[0]\n",
    "print(f\"{dataset} dataset with {num_training_data} training samples and {num_testing_data} testing samples\"\n",
    "      f\" and {input_dim} features and \"\n",
    "      f\"{torch.unique(torch.tensor(loader_train.dataset.targets), return_counts=True)} classes.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This mlp has 21.302 thousand parameters.\n"
     ]
    }
   ],
   "source": [
    "# create a model trained with pear (lambda = 0.5)\n",
    "disagreement_lambda = 0.5\n",
    "epochs = 50\n",
    "\n",
    "model_pear = pear.get_model(model_cfg, input_dim, num_classes)\n",
    "pytorch_total_params = sum(p.numel() for p in model_pear.parameters())\n",
    "print(f\"This {model_cfg['name']} has {pytorch_total_params / 1e3:0.3f} thousand parameters.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get an optimizer\n",
    "params = model_pear.parameters()\n",
    "optim = AdamW(params, lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "# create two explainers for the loss\n",
    "explainer_a = pear.get_explainer(explainers[0], model_pear, torch.tensor(loader_train.dataset.data))\n",
    "explainer_b = pear.get_explainer(explainers[1], model_pear, torch.tensor(loader_train.dataset.data))\n",
    "disagreement_loss_fn = pear.DisagreementLoss(explainer_a, explainer_b, disagreement_mu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  0 | task loss 0.4462 | disagree loss 0.0035 | train bal acc 78.04 | test bal acc 78.80 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  1 | task loss 0.4118 | disagree loss 0.0034 | train bal acc 80.11 | test bal acc 80.73 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  2 | task loss 0.4014 | disagree loss 0.0028 | train bal acc 80.90 | test bal acc 81.37 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  3 | task loss 0.4015 | disagree loss 0.0030 | train bal acc 80.87 | test bal acc 80.28 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  4 | task loss 0.3946 | disagree loss 0.0034 | train bal acc 81.39 | test bal acc 81.65 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  5 | task loss 0.4092 | disagree loss 0.0024 | train bal acc 80.65 | test bal acc 81.37 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  6 | task loss 0.3920 | disagree loss 0.0029 | train bal acc 81.39 | test bal acc 81.52 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  7 | task loss 0.3903 | disagree loss 0.0027 | train bal acc 81.50 | test bal acc 82.12 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  8 | task loss 0.4034 | disagree loss 0.0038 | train bal acc 80.65 | test bal acc 79.92 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  9 | task loss 0.3919 | disagree loss 0.0032 | train bal acc 81.58 | test bal acc 82.01 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 10 | task loss 0.3876 | disagree loss 0.0030 | train bal acc 81.80 | test bal acc 82.00 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 11 | task loss 0.3899 | disagree loss 0.0029 | train bal acc 81.72 | test bal acc 81.59 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 12 | task loss 0.3850 | disagree loss 0.0036 | train bal acc 81.97 | test bal acc 82.20 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 13 | task loss 0.3875 | disagree loss 0.0046 | train bal acc 81.88 | test bal acc 81.68 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 14 | task loss 0.3836 | disagree loss 0.0041 | train bal acc 82.17 | test bal acc 82.33 | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▉         | 22/242 [00:00<00:02, 102.29it/s]"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    _ = model_pear.train_loop(trainloader=loader_train,\n",
    "                              disagreement_lambda=disagreement_lambda,\n",
    "                              optimizer=optim,\n",
    "                              task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                              disagreement_loss_fn=disagreement_loss_fn)\n",
    "    evaluation_on_train_data = model_pear.evaluate_balanced(loader_train,\n",
    "                                                            task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                                            disagreement_loss_fn=disagreement_loss_fn)\n",
    "    evaluation_on_test_data = model_pear.evaluate_balanced(loader_test,\n",
    "                                                           task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                                           disagreement_loss_fn=disagreement_loss_fn)\n",
    "\n",
    "    print(f\"epoch {epoch:2d} | \"\n",
    "          f\"task loss {evaluation_on_train_data['task_loss']:.4f} | \"\n",
    "          f\"disagree loss {evaluation_on_train_data['disagreement_loss']:.4f} | \"\n",
    "          f\"train bal acc {(evaluation_on_train_data['acc_0'] + evaluation_on_train_data['acc_1']) / 2:.2f} | \"\n",
    "          f\"test bal acc {(evaluation_on_test_data['acc_0'] + evaluation_on_test_data['acc_1']) / 2:.2f} | \"\n",
    "          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a second model with a different lambda\n",
    "disagreement_lambda = 0.0\n",
    "epochs = 30\n",
    "\n",
    "model_vanilla = pear.get_model(model_cfg, input_dim, num_classes)\n",
    "pytorch_total_params = sum(p.numel() for p in model_vanilla.parameters())\n",
    "print(f\"This {model_cfg['name']} has {pytorch_total_params / 1e3:0.3f} thousand parameters.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get an optimizer\n",
    "params = model_vanilla.parameters()\n",
    "optim = AdamW(params, lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "# create two explainers for the loss\n",
    "explainer_a = pear.get_explainer(explainers[0], model_vanilla, torch.tensor(loader_train.dataset.data))\n",
    "explainer_b = pear.get_explainer(explainers[1], model_vanilla, torch.tensor(loader_train.dataset.data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    _ = model_vanilla.train_loop(trainloader=loader_train,\n",
    "                                 disagreement_lambda=disagreement_lambda,\n",
    "                                 optimizer=optim,\n",
    "                                 task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                 disagreement_loss_fn=disagreement_loss_fn)\n",
    "    evaluation_on_train_data = model_vanilla.evaluate_balanced(loader_train,\n",
    "                                                               task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                                               disagreement_loss_fn=disagreement_loss_fn)\n",
    "    evaluation_on_test_data = model_vanilla.evaluate_balanced(loader_test,\n",
    "                                                              task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                                              disagreement_loss_fn=disagreement_loss_fn)\n",
    "\n",
    "    print(f\"epoch {epoch:2d} | \"\n",
    "          f\"task loss {evaluation_on_train_data['task_loss']:.4f} | \"\n",
    "          f\"disagree loss {evaluation_on_train_data['disagreement_loss']:.4f} | \"\n",
    "          f\"train bal acc {(evaluation_on_train_data['acc_0'] + evaluation_on_train_data['acc_1']) / 2:.2f} | \"\n",
    "          f\"test bal acc {(evaluation_on_test_data['acc_0'] + evaluation_on_test_data['acc_1']) / 2:.2f} | \"\n",
    "          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "metric = \"pairwise_rank\"\n",
    "red_grid_data_pear = pear.disagreement_matrices(model_pear, loader_train, loader_test, k=5, metric=metric)\n",
    "red_grid_data_vanilla = pear.disagreement_matrices(model_vanilla, loader_train, loader_test, k=5, metric=metric)\n",
    "\n",
    "explainer_indices = {\"vanilla_gradients\": 2,\n",
    "                     \"integrated_gradients\": 4,\n",
    "                     \"shap\": 1,\n",
    "                     \"lime\": 0,\n",
    "                     \"input_x_gradient\": 3,\n",
    "                     \"smooth_grad\": 5}\n",
    "\n",
    "explainer_pairs = [\n",
    "    \"input_x_gradient_v_input_x_gradient\",\n",
    "    \"input_x_gradient_v_integrated_gradients\",\n",
    "    \"input_x_gradient_v_lime\",\n",
    "    \"input_x_gradient_v_shap\",\n",
    "    \"input_x_gradient_v_smooth_grad\",\n",
    "    \"input_x_gradient_v_vanilla_gradients\",\n",
    "    \"integrated_gradients_v_integrated_gradients\",\n",
    "    \"integrated_gradients_v_lime\",\n",
    "    \"integrated_gradients_v_shap\",\n",
    "    \"integrated_gradients_v_smooth_grad\",\n",
    "    \"integrated_gradients_v_vanilla_gradients\",\n",
    "    \"lime_v_lime\",\n",
    "    \"lime_v_shap\",\n",
    "    \"lime_v_smooth_grad\",\n",
    "    \"lime_v_vanilla_gradients\",\n",
    "    \"shap_v_shap\",\n",
    "    \"shap_v_smooth_grad\",\n",
    "    \"shap_v_vanilla_gradients\",\n",
    "    \"smooth_grad_v_smooth_grad\",\n",
    "    \"smooth_grad_v_vanilla_gradients\",\n",
    "    \"vanilla_gradients_v_vanilla_gradients\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_strs = {\n",
    "    \"feature_agreement\": \"Feature Agreement\",\n",
    "    \"rank_agreement\": \"Rank Agreement\",\n",
    "    \"sign_agreement\": \"Sign Agreement\",\n",
    "    \"signed_rank_agreement\": \"Signed Rank Agreement\",\n",
    "    \"rank_correlation\": \"Rank Correlation\",\n",
    "    \"pairwise_rank\": \"Pairwise Rank Agreement\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "cmap = sns.color_palette(\"light:darkred\", as_cmap=True)\n",
    "metric = \"pairwise_rank\"\n",
    "fs = 12\n",
    "tables = [red_grid_data_pear, red_grid_data_vanilla]\n",
    "lams = [0.5, 0.0]\n",
    "for table, lam in zip(tables, lams):\n",
    "    to_plot = np.zeros((len(explainer_indices.keys()), len(explainer_indices.keys())))\n",
    "    for explainer_pair in explainer_pairs:\n",
    "        exs = explainer_pair.split(\"_v_\")\n",
    "        to_plot[explainer_indices[exs[0]], explainer_indices[exs[1]]] = table[metric][explainer_pair]\n",
    "        to_plot[explainer_indices[exs[1]], explainer_indices[exs[0]]] = table[metric][explainer_pair]\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(5, 4.8))\n",
    "    if \"correlation\" in metric:\n",
    "        sns.heatmap(to_plot, vmin=-1, vmax=1, cmap=cmap, ax=ax, annot=True)\n",
    "        title = f\"California Housing Data\\n{metric_strs[metric]}\\n$\\lambda$ = {lam}\"\n",
    "    else:\n",
    "        sns.heatmap(to_plot, vmin=0, vmax=1, cmap=cmap, ax=ax, annot=True)\n",
    "        title = f\"California Housing Data\\n{metric_strs[metric]}\\n$\\lambda$ = {lam} and k = {5}\"\n",
    "    ax.set_title(title, fontsize=fs)\n",
    "\n",
    "    ax.set_xticks([0 + 0.5, 1 + 0.5, 2 + 0.5, 3 + 0.5, 4 + 0.5, 5 + 0.5],\n",
    "                  [\"LIME\", \"SHAP\", \"Grad\", \"Grad*\\nInput\", \"IntGrad\", \"Smooth\\nGrad\"],\n",
    "                  rotation=0,\n",
    "                  fontsize=fs)\n",
    "    ax.set_yticks([0 + 0.5, 1 + 0.5, 2 + 0.5, 3 + 0.5, 4 + 0.5, 5 + 0.5],\n",
    "                  [\"LIME\", \"SHAP\", \"Grad\", \"Grad*\\nInput\", \"IntGrad\", \"Smooth\\nGrad\"],\n",
    "                  rotation=0,\n",
    "                  fontsize=fs)\n",
    "    plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(12, 9))\n",
    "title = f\"California Housing Data\\nLIME vs. SHAP\"\n",
    "fs = 48\n",
    "ax.set_title(title, fontsize=fs)\n",
    "\n",
    "vanilla_acc = model_vanilla.evaluate(loader_test,\n",
    "                                     task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                     disagreement_loss_fn=disagreement_loss_fn)[\"acc\"]\n",
    "vanilla_agreement = red_grid_data_vanilla[metric][\"lime_v_shap\"]\n",
    "ax.plot([vanilla_acc], [vanilla_agreement],\n",
    "        marker=\"*\",\n",
    "        markersize=22,\n",
    "        linestyle=\"none\",\n",
    "        color=\"grey\",\n",
    "        label=\"Vanilla\")\n",
    "\n",
    "pear_acc = model_vanilla.evaluate(loader_test,\n",
    "                                  task_loss_fn=torch.nn.CrossEntropyLoss(),\n",
    "                                  disagreement_loss_fn=disagreement_loss_fn)[\"acc\"]\n",
    "pear_agreement = red_grid_data_pear[metric][\"lime_v_shap\"]\n",
    "ax.plot([pear_acc], [pear_agreement],\n",
    "        marker=\"o\",\n",
    "        markersize=22,\n",
    "        linestyle=\"none\",\n",
    "        color=\"red\",\n",
    "        label=\"PEAR\")\n",
    "\n",
    "ax.set_xlim([81, 85.5])\n",
    "x = np.arange(81, 86, 1)\n",
    "ax.set_xticks(x)\n",
    "ax.set_xticklabels(x, fontsize=fs, rotation=0)\n",
    "ax.tick_params(axis='x', labelsize=fs - 8)\n",
    "ax.set_xlabel(\"Test Accuracy (%)\", fontsize=fs)\n",
    "\n",
    "ax.set_ylim([0.7, 0.85])\n",
    "y = np.arange(0.72, 0.85, 0.04)\n",
    "ax.set_yticks(y)\n",
    "ax.set_yticklabels([f\"{i:0.2f}\" for i in y], fontsize=fs, rotation=0)\n",
    "ax.tick_params(axis='y', labelsize=fs - 8)\n",
    "ax.set_ylabel(f\"{metric_strs[metric]}\", fontsize=fs)\n",
    "\n",
    "ax.spines[\"right\"].set_visible(False)\n",
    "ax.spines[\"top\"].set_visible(False)\n",
    "ax.legend(loc='upper left', bbox_to_anchor=(0.2, 0.5), fontsize=fs / 2)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}