Exp2_Hessian_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2aa6ffd9",
   "metadata": {},
   "source": [
    "## Import Package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "81e9dc52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11/11 08:56:01 AM Note: detected 78 virtual cores but NumExpr set to maximum of 64, check \"NUMEXPR_MAX_THREADS\" environment variable.\n",
      "11/11 08:56:01 AM Note: NumExpr detected 78 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error.  nthreads cannot be larger than environment variable \"NUMEXPR_MAX_THREADS\" (64)"
     ]
    }
   ],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import pprint\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import pickle\n",
    "import copy\n",
    "import collections\n",
    "import math\n",
    "\n",
    "import numpy as np\n",
    "import numpy\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler,TensorDataset\n",
    "\n",
    "from torch.nn import CrossEntropyLoss, MSELoss\n",
    "\n",
    "from transformer import BertForSequenceClassification,WEIGHTS_NAME, CONFIG_NAME\n",
    "from transformer.modeling_quant import BertForSequenceClassification as QuantBertForSequenceClassification\n",
    "from transformer import BertTokenizer\n",
    "from transformer import BertAdam\n",
    "from transformer import BertConfig\n",
    "from transformer import QuantizeLinear, BertSelfAttention, FP_BertSelfAttention, BertAttention, FP_BertAttention\n",
    "from utils_glue import *\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import torch.nn.functional as F\n",
    "\n",
    "mse_func = MSELoss()\n",
    "\n",
    "processors = {\n",
    "    \"cola\": ColaProcessor,\n",
    "    \"mnli\": MnliProcessor,\n",
    "    \"mnli-mm\": MnliMismatchedProcessor,\n",
    "    \"mrpc\": MrpcProcessor,\n",
    "    \"sst-2\": Sst2Processor,\n",
    "    \"sts-b\": StsbProcessor,\n",
    "    \"qqp\": QqpProcessor,\n",
    "    \"qnli\": QnliProcessor,\n",
    "    \"rte\": RteProcessor   \n",
    "}\n",
    "\n",
    "output_modes = {\n",
    "        \"cola\": \"classification\",\n",
    "        \"mnli\": \"classification\",\n",
    "        \"mrpc\": \"classification\",\n",
    "        \"sst-2\": \"classification\",\n",
    "        \"sts-b\": \"regression\",\n",
    "        \"qqp\": \"classification\",\n",
    "        \"qnli\": \"classification\",\n",
    "        \"rte\": \"classification\"\n",
    "}\n",
    "\n",
    "default_params = {\n",
    "        \"cola\": {\"max_seq_length\": 64,\"batch_size\":1,\"eval_step\": 50}, # No Aug : 50 Aug : 400\n",
    "        \"mnli\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\":8000},\n",
    "        \"mrpc\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\":100},\n",
    "        \"sst-2\": {\"max_seq_length\": 64,\"batch_size\":1,\"eval_step\":100},\n",
    "        \"sts-b\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\":100},\n",
    "        \"qqp\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\":1000},\n",
    "        \"qnli\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\":1000},\n",
    "        \"rte\": {\"max_seq_length\": 128,\"batch_size\":1,\"eval_step\": 20}\n",
    "    }\n",
    "\n",
    "def get_tensor_data(output_mode, features):\n",
    "    if output_mode == \"classification\":\n",
    "        all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.long)\n",
    "    elif output_mode == \"regression\":\n",
    "        all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.float)\n",
    "\n",
    "\n",
    "    all_seq_lengths = torch.tensor([f.seq_length for f in features], dtype=torch.long)\n",
    "    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)\n",
    "    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)\n",
    "    tensor_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,all_label_ids, all_seq_lengths)\n",
    "    return tensor_data, all_label_ids\n",
    "\n",
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0 \n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88608e80",
   "metadata": {},
   "source": [
    "## Dataset & Model Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d911623a",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_name = \"sst-2\"\n",
    "bert_size = \"large\"\n",
    "\n",
    "if bert_size == \"large\":\n",
    "    layer_num = 24\n",
    "    head_num = 16\n",
    "else: \n",
    "    layer_num = 12\n",
    "    head_num = 12\n",
    "\n",
    "model_dir = \"models/BERT_large/sst-2\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad43e12a",
   "metadata": {},
   "source": [
    "## Prepare Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "35278bbc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11/11 08:56:01 AM Writing example 0 of 872\n",
      "11/11 08:56:01 AM *** Example ***\n",
      "11/11 08:56:01 AM guid: dev-1\n",
      "11/11 08:56:01 AM tokens: [CLS] it ' s a charming and often affecting journey . [SEP]\n",
      "11/11 08:56:01 AM input_ids: 101 2009 1005 1055 1037 11951 1998 2411 12473 4990 1012 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "11/11 08:56:01 AM input_mask: 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "11/11 08:56:01 AM segment_ids: 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "11/11 08:56:01 AM label: 1\n",
      "11/11 08:56:01 AM label_id: 1\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Processor & Task Info\n",
    "processor = processors[task_name]()\n",
    "output_mode = output_modes[task_name]\n",
    "label_list = processor.get_labels()\n",
    "num_labels = len(label_list)\n",
    "\n",
    "if task_name in default_params:\n",
    "    batch_size = default_params[task_name][\"batch_size\"]\n",
    "    max_seq_length = default_params[task_name][\"max_seq_length\"]\n",
    "    eval_step = default_params[task_name][\"eval_step\"]\n",
    "    \n",
    "# Tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained(model_dir, do_lower_case=True)\n",
    "\n",
    "# Load Dataset\n",
    "data_dir = os.path.join(\"data\",task_name)\n",
    "\n",
    "eval_examples = processor.get_dev_examples(data_dir)\n",
    "eval_features = convert_examples_to_features(eval_examples, label_list, max_seq_length, tokenizer, output_mode)\n",
    "\n",
    "eval_data, eval_labels = get_tensor_data(\"classification\", eval_features)\n",
    "eval_sampler = RandomSampler(eval_data)\n",
    "eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=1)\n",
    "eval_data, eval_labels = get_tensor_data(output_mode, eval_features)\n",
    "\n",
    "# Get input batch sample\n",
    "batch = next(iter(eval_dataloader))\n",
    "input_ids, input_mask, segment_ids, label_ids, seq_lengths = batch\n",
    "seq_length = seq_lengths[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b58137df",
   "metadata": {},
   "source": [
    "## Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a46c351",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Student Model - Attention Map KD-QAT\n"
     ]
    }
   ],
   "source": [
    "# Build Model\n",
    "student_model_M_dir = \"output/BERT_large/sst-2/exploration/sst-2_large_map_large_M_42\"\n",
    "student_config = BertConfig.from_pretrained(student_model_M_dir)             \n",
    "student_model_M = QuantBertForSequenceClassification.from_pretrained(student_model_M_dir, config = student_config, num_labels=num_labels)\n",
    "print(\"Student Model - Attention Map KD-QAT\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eed30a9",
   "metadata": {},
   "source": [
    "## Hessian Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07a1df63",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hessian import hessian\n",
    "\n",
    "if output_mode == \"classification\":\n",
    "    loss_fct = CrossEntropyLoss()\n",
    "elif output_mode == \"regression\":\n",
    "    loss_fct = MSELoss()\n",
    "\n",
    "tc_max_eigens = []\n",
    "model = student_model_M.to(device)\n",
    "teacher_model = None\n",
    "\n",
    "for batch in tqdm(eval_dataloader):\n",
    "    input_ids, input_mask, segment_ids, label_ids, seq_lengths = batch\n",
    "    hessian_comp = hessian(model, data=(input_ids, label_ids), criterion=loss_fct, cuda=True, input_zip = (input_ids, segment_ids, input_mask), teacher_model=teacher_model, kd_type=None)\n",
    "    top_eigenvalues, top_eigenvector = hessian_comp.eigenvalues(top_n=3)\n",
    "    tc_max_eigens = tc_max_eigens + top_eigenvalues"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f5a6a80",
   "metadata": {},
   "source": [
    "## Save Hessian Max Eigen Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6e379a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_folder_name = \"hessian_value_pts\"\n",
    "file_name = f\"{bert_size}_{task_name}_{hessian}.pt\"}\n",
    "\n",
    "if not os.path.exists(pt_folder_name):\n",
    "    os.mkdir(pt_folder_name)\n",
    "    \n",
    "file_dir = os.path.join(pt_folder_name, file_name)\n",
    "\n",
    "print(file_dir)\n",
    "try:\n",
    "    torch.save(tc_max_eigens, file_dir)\n",
    "except:\n",
    "    import pdb; pdb.set_trace()\n",
    "print(\"==> Model Eigen Value DONE!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad8c870",
   "metadata": {},
   "source": [
    "## Plot Hessian Eigenvalue Spectra (Figure. 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a945b563",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "eigens_1 = torch.load(\"hessian_value_pts/cola_sarq_init.pt\")\n",
    "eigens_2 =torch.load(\"hessian_value_pts/cola_ternary_init.pt\")\n",
    "\n",
    "eigens_1_pos = []\n",
    "eigens_1_neg = []\n",
    "for eigen in eigens_1:\n",
    "    if eigen > 0:\n",
    "        eigens_1_pos.append(eigen)\n",
    "    else:\n",
    "        eigens_1_neg.append(eigen)\n",
    "    \n",
    "eigens_2_pos = []\n",
    "eigens_2_neg = []\n",
    "for eigen in eigens_2:\n",
    "    if eigen > 0:\n",
    "        eigens_2_pos.append(eigen)\n",
    "    else:\n",
    "        eigens_2_neg.append(eigen)\n",
    "        \n",
    "fs = 13\n",
    "lw = 2.5\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10,3.4), dpi=200)\n",
    "\n",
    "color_1 = \"tab:blue\"\n",
    "color_2 = \"navy\"\n",
    "color_3 = \"darkblue\"\n",
    "color_4 = \"tab:red\"\n",
    "\n",
    "# pos = [pos_1, pos_2]\n",
    "# Plot\n",
    "pos_1 = sns.kdeplot(eigens_1_pos, color=color_1, label=eigens_1_name, linewidth=lw, ax=axes[1])\n",
    "pos_2 = sns.kdeplot(eigens_2_pos, color=color_2, label=eigens_2_name, linewidth=lw, ax=axes[1])\n",
    "neg_1 = sns.kdeplot(eigens_1_neg, color=color_1, label=eigens_1_name_2, linewidth=lw, ax=axes[0])\n",
    "neg_2 = sns.kdeplot(eigens_2_neg, color=color_2, label=eigens_2_name_2, linewidth=lw, ax=axes[0])\n",
    "\n",
    "# Font Size\n",
    "\n",
    "pos_1.get_yaxis().set_visible(False)\n",
    "pos_1.tick_params(axis='x', labelsize=fs)\n",
    "neg_1.tick_params(axis='x', labelsize=fs)\n",
    "neg_1.tick_params(axis='y', labelsize=fs)\n",
    "neg_1.set_ylabel(ylabel=\"Density\", fontsize = fs+2)\n",
    "pos_1.set_xlabel(xlabel=\"Positive Max Eigenvalue\", fontsize = fs+2)\n",
    "neg_1.set_xlabel(xlabel=\"Negative Max Eigenvalue\", fontsize = fs+2)\n",
    "axes[1].legend(fontsize = fs, loc=1)\n",
    "axes[0].legend(fontsize = fs, loc=2)\n",
    "axes[1].set_ylim(0, 0.025)\n",
    "fig.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch 1.13 (NGC 22.05/Python 3.8 Conda) on Backend.AI",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}