notebooks/plot_classifier_weights_biases.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e7e41b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Plot analysis of classifier weight and bias norms.\n",
    "\n",
    "Need to launch with .sh script importing PYTHON_PATH for loading ckpt.\n",
    "Before running this experiment, make sure you have run the Stream meta-data collector (src/continual_ego4d/processing/run_summarize_user_streams.py) and set the resulting paths in this notebook.\n",
    "\"\"\"\n",
    "import datetime\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "from collections import Counter\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import scipy\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "\n",
    "InteractiveShell.ast_node_interactivity = \"all\"  # Print all variables on their own lines\n",
    "pd.set_option('display.max_rows', 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "CONFIG: Add your config params here\n",
    "\"\"\"\n",
    "\n",
    "########### PLOT CONFIG ###############\n",
    "plot_modes = ['verb', 'noun'] # Make plot for verbs or for nouns classifier\n",
    "plot_mode = plot_modes[0]\n",
    "\n",
    "########### DATA CONFIG ###############\n",
    "# Path obtained by Stream meta-data collector (src/continual_ego4d/processing/run_summarize_user_streams.py)\n",
    "train_user_stream_summary_path = '/your/path/to/logs/2022-09-12_18-25-12_UID6131a811-f2c5-4479-a7ff-08dc74d4f9fc/dataset_entries_train_FEWSHOT=False_ego4d_LTA_train_usersplit_10users.ckpt' # Excludes unseen actions during pretrain\n",
    "\n",
    "# Path to direct JSON from pretraining\n",
    "pretrain_path = '../data/EgoAdapt/usersplits/ego4d_LTA_pretrain_incl_nanusers_usersplit_148users.json'\n",
    "\n",
    "########### MODELS CONFIG ###############\n",
    "# SELECT MODEL TO ANALYZE\n",
    "\n",
    "# SGD training full model\n",
    "sgd_ckp_model_dirs = '/your/path/to/logs/GRID_SOLVER-BASE_LR=0-01_SOLVER-MOMENTUM=0-0_SOLVER-NESTEROV=True/2022-09-13_10-53-52_UID958392f7-c477-4a09-a7ac-c72cc81251c2/checkpoints'\n",
    "\n",
    "# Replay full storage, lr 0.01\n",
    "replay_ckp_model_dirs = '/your/path/to/logs/GRID_METHOD-REPLAY-MEMORY_SIZE_SAMPLES=1000000_METHOD-REPLAY-STORAGE_POLICY=window/2022-09-15_18-23-13_UID09ab4f67-814b-4331-aa80-839cd99a1d9f/checkpoints'\n",
    "\n",
    "# SGD training classifier only:\n",
    "sgd_classifier_only_ckp_model_dirs = \"/your/path/to/logs/GRID_SOLVER-BASE_LR=0-1/2022-09-19_12-16-57_UID11a32c82-471b-4921-a5c7-0ee393fdabf0/checkpoints\"\n",
    "\n",
    "# Pretrained model\n",
    "pretrain_model_path = '/your/path/to/pretrain_148usersplit_incl_nan/2022-09-05_10-34-05_UIDd05ed672-01c5-4c3c-b790-9d0c76548825/checkpoints/best_model.ckpt'"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7656eb7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    " Get all seen verbs/nouns/actions in pretrain.\n",
    "\"\"\"\n",
    "with open(pretrain_path, \"r\") as f:\n",
    "    pretrain_dataset = json.load(f)\n",
    "\n",
    "pretrain_action_sets = pretrain_dataset['user_action_sets']['user_agnostic']\n",
    "\n",
    "pretrain_verb_freq_dict = {\n",
    "    int(a): a_dict['count'] for a, a_dict in pretrain_action_sets['verb_to_name_dict'].items()\n",
    "}\n",
    "pretrain_noun_freq_dict = {\n",
    "    int(a): a_dict['count'] for a, a_dict in pretrain_action_sets['noun_to_name_dict'].items()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9fbe8df",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Get dictionary with dataframe per user stream (includes all labels and other meta data).\"\"\"\n",
    "\n",
    "def ds_to_user_dfs(ds):\n",
    "    ret = {}\n",
    "    for user, user_entries in ds.items():\n",
    "        # Do all for actions/verbs/nouns\n",
    "        user_df = pd.json_normalize(user_entries)  # Convert to DF\n",
    "\n",
    "        # Create action column\n",
    "        def label_fn(x):\n",
    "            assert len(x) == 2, \"Need two columns to merge\"\n",
    "            if not isinstance(x[0], list):\n",
    "                assert not isinstance(x[1], list)\n",
    "                return f\"{x[0]}-{x[1]}\"\n",
    "\n",
    "            return [f\"{l}-{r}\" for l, r in zip(x[0], x[1])]\n",
    "\n",
    "        user_df['action_label'] = user_df.loc[:, ('verb_label', 'noun_label')].apply(label_fn, axis=1)\n",
    "        ret[user] = user_df\n",
    "        print(\"Created action_label column\")\n",
    "    return ret\n",
    "\n",
    "\n",
    "with open(train_user_stream_summary_path, 'rb') as f:\n",
    "    ds = pickle.load(f)\n",
    "dfs = ds_to_user_dfs(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f4de1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "COLLECT ALL WEIGHTS AND BIASES for verb or noun (select which one).\n",
    "These stats are later used for plotting.\n",
    "\"\"\"\n",
    "head_layers = [\n",
    "    'model.head.0.projections.0.weight',\n",
    "    'model.head.0.projections.0.bias',\n",
    "    'model.head.0.projections.1.weight',\n",
    "    'model.head.0.projections.1.bias']\n",
    "\n",
    "verb_bias_name = 'model.head.0.projections.0.bias'\n",
    "verb_weight_name = 'model.head.0.projections.0.weight'\n",
    "noun_bias_name = 'model.head.0.projections.1.bias'\n",
    "noun_weight_name = 'model.head.0.projections.1.weight'\n",
    "\n",
    "verb_idx_to_keep = list(pretrain_verb_freq_dict.keys())\n",
    "noun_idx_to_keep = list(pretrain_noun_freq_dict.keys())\n",
    "\n",
    "max_verbs_total = 115\n",
    "max_nouns_total = 478\n",
    "print(f\"VERBS: Keep {len(verb_idx_to_keep)}/{max_verbs_total}\")\n",
    "print(f\"NOUNS: Keep {len(noun_idx_to_keep)}/{max_nouns_total}\")\n",
    "\n",
    "if plot_mode == 'verb':\n",
    "    bias_name = verb_bias_name\n",
    "    weight_name = verb_weight_name\n",
    "    idx_to_keep = verb_idx_to_keep\n",
    "\n",
    "    pretrain_bias_name = 'model.head.projections.0.bias'\n",
    "    pretrain_weight_name = 'model.head.projections.0.weight'\n",
    "    stream_stats_key = 'verb_label'\n",
    "    max_total = max_verbs_total\n",
    "\n",
    "elif plot_mode == 'noun':\n",
    "    bias_name = noun_bias_name\n",
    "    weight_name = noun_weight_name\n",
    "    idx_to_keep = noun_idx_to_keep\n",
    "\n",
    "    pretrain_bias_name = 'model.head.projections.1.bias'\n",
    "    pretrain_weight_name = 'model.head.projections.1.weight'\n",
    "    stream_stats_key = 'noun_label'\n",
    "    max_total = max_nouns_total\n",
    "\n",
    "\n",
    "def get_bias_selection_distr(model_state_dict, verb_bias_name, verb_idx_to_keep, full_bias_t=None):\n",
    "    if full_bias_t is None:\n",
    "        full_bias_t = model_state_dict[verb_bias_name]\n",
    "\n",
    "    verb_bias_t = full_bias_t[verb_idx_to_keep]  # Drop unseens in pretrain\n",
    "    verb_bias_t_pos = verb_bias_t.pow(2).sqrt()  # Squared L2\n",
    "    verb_bias_tn = (verb_bias_t_pos) / torch.sum(verb_bias_t_pos)  # Normalize to distr\n",
    "    verb_bias_list = verb_bias_tn.tolist()\n",
    "    return verb_bias_list\n",
    "\n",
    "\n",
    "def get_weight_selection_distr(model_state_dict, verb_weight_name, verb_idx_to_keep, full_weight_t=None):\n",
    "    if full_weight_t is None:\n",
    "        full_weight_t = model_state_dict[verb_weight_name]  # torch.Size([115, 2304])\n",
    "\n",
    "    verb_weight_t = full_weight_t[verb_idx_to_keep]  # torch.Size([106, 2304])\n",
    "    verb_weight_norms_t = verb_weight_t.pow(2).sum(dim=1)  # All >= 0\n",
    "    verb_weight_norms_tn = verb_weight_norms_t / torch.sum(verb_weight_norms_t)\n",
    "    verb_weight_list = verb_weight_norms_tn.tolist()\n",
    "    return verb_weight_list\n",
    "\n",
    "\n",
    "# GET PRETRAINED\n",
    "ckpt = torch.load(pretrain_model_path, map_location='cpu')\n",
    "pretrain_state_dict = ckpt['state_dict']\n",
    "\n",
    "verb_bias_distr_pretrain = get_bias_selection_distr(pretrain_state_dict, pretrain_bias_name, idx_to_keep)\n",
    "verb_weight_distr_pretrain = get_weight_selection_distr(pretrain_state_dict, pretrain_weight_name, idx_to_keep)\n",
    "# For pretrain_state_dict.keys(): 'model.head.projections.0.weight', 'model.head.projections.0.bias', 'model.head.projections.1.weight', 'model.head.projections.1.bias']\n",
    "\n",
    "# GET USERS\n",
    "model_results = {\n",
    "    'sgd': {'path': sgd_ckp_model_dirs},\n",
    "    'replay': {'path': replay_ckp_model_dirs},\n",
    "    'sgd_classifier_only': {'path': sgd_classifier_only_ckp_model_dirs},\n",
    "}\n",
    "\n",
    "for name, model_result in model_results.items():\n",
    "    print(f\"Processing: {name}\")\n",
    "\n",
    "    df_list: list[dict] = []\n",
    "    for idx, (user_id, user_df) in enumerate(dfs.items()):\n",
    "        print(f\"Processing user-id:{user_id}\")\n",
    "\n",
    "        # STREAM STATS\n",
    "        cnt_dict = Counter(user_df[stream_stats_key].tolist())\n",
    "        verb_cnt_t = torch.zeros(max_total)  # Original idxs\n",
    "\n",
    "        for verb_idx, cnt in cnt_dict.items():\n",
    "            verb_cnt_t[verb_idx] = cnt\n",
    "\n",
    "        verb_cnt_t = verb_cnt_t[idx_to_keep]  # Subset idxs\n",
    "        verb_cnt = verb_cnt_t.tolist()\n",
    "        verb_distr_stream = (verb_cnt_t / verb_cnt_t.abs().sum()).tolist()  # Divide by sum\n",
    "\n",
    "        # MODEL \n",
    "        user_path = os.path.join(model_result['path'], f\"user_{user_id}\", 'last.ckpt')\n",
    "        ckpt = torch.load(user_path, map_location='cpu')\n",
    "        model_state_dict = ckpt['state_dict']\n",
    "\n",
    "        # GET DELTAS w.r.t. Pretrain\n",
    "        # Biases\n",
    "        bias_deltas = (model_state_dict[bias_name] - pretrain_state_dict[pretrain_bias_name])[idx_to_keep]\n",
    "        bias_deltas_n = (bias_deltas / bias_deltas.abs().sum()).tolist()  # Normalize by absolute mass\n",
    "\n",
    "        # Weights\n",
    "        weight_norm_deltas = (model_state_dict[weight_name].pow(2).sum(dim=1).sqrt() - pretrain_state_dict[\n",
    "            pretrain_weight_name].pow(2).sum(dim=1).sqrt())[idx_to_keep]\n",
    "        weight_norm_deltas_n = (weight_norm_deltas / weight_norm_deltas.abs().sum()).tolist()\n",
    "\n",
    "        for idx in range(len(idx_to_keep)):\n",
    "            df_list.append({\n",
    "                f\"id_{plot_mode}\": idx,\n",
    "                \"user\": user_id,\n",
    "                'stream_cnt': verb_cnt[idx],\n",
    "                # Distributions:\n",
    "                \"weight_norm_delta_p\": weight_norm_deltas_n[idx],\n",
    "                \"bias_delta_p\": bias_deltas_n[idx],\n",
    "                'stream_p': verb_distr_stream[idx],\n",
    "            })\n",
    "\n",
    "    # Store results for sgd/replay\n",
    "    users_df = pd.DataFrame(df_list)\n",
    "    model_result['df'] = users_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a40d3e36",
   "metadata": {
    "scrolled": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Plot\n",
    "\"\"\"\n",
    "\n",
    "ZOOM = True\n",
    "xlim = (-2, None)\n",
    "PLOT_SELECTION = ['sgd', 'sgd_classifier_only']  # 'replay'\n",
    "\n",
    "legend_name_mapping = {\n",
    "    'sgd': r'SGD - $F \\circ H$',\n",
    "    'sgd_classifier_only': r'SGD - $H$',\n",
    "    'replay': 'replay',\n",
    "}\n",
    "\n",
    "# Set fonts\n",
    "plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "plt.rcParams['font.size'] = 14\n",
    "\n",
    "# Use latex in mpl\n",
    "plt.rcParams['text.usetex'] = True\n",
    "plt.rcParams['text.latex.preamble'] = r'\\usepackage{amsmath}'  #for \\text command\n",
    "\n",
    "# Set latex font in mpl\n",
    "plt.rcParams['mathtext.fontset'] = 'custom'\n",
    "plt.rcParams['mathtext.rm'] = 'Bitstream Vera Sans'\n",
    "plt.rcParams['mathtext.it'] = 'Bitstream Vera Sans:italic'\n",
    "plt.rcParams['mathtext.bf'] = 'Bitstream Vera Sans:bold'\n",
    "\n",
    "figsize = (5, 4)\n",
    "\n",
    "# Paths\n",
    "for plot_weights in [True, False]:\n",
    "    main_outdir = \"../imgs/classifier_analysis_final\"\n",
    "    title = f\"classifier_analysis_{plot_mode}_\"\n",
    "    title += \"weight\" if plot_weights else \"bias\"\n",
    "    parent_dirname = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + \"_\" + title\n",
    "    parent_dirpath = os.path.join(main_outdir, parent_dirname)\n",
    "    filename = f\"{title}.pdf\"\n",
    "    os.makedirs(parent_dirpath, exist_ok=True)\n",
    "\n",
    "    # PLOT\n",
    "    fig = plt.figure(figsize=figsize, dpi=600)  # So all bars are visible!\n",
    "    ax = plt.subplot()\n",
    "    labels = []\n",
    "\n",
    "    # Plot histogram distr\n",
    "    sgd_df = model_results['sgd']['df']\n",
    "    df_group = sgd_df.groupby([f\"id_{plot_mode}\"], as_index=False).agg({'stream_p': 'mean'})\n",
    "\n",
    "    sort_idxs = list(reversed(df_group['stream_p'].argsort()))\n",
    "    df_group = df_group.loc[sort_idxs]\n",
    "\n",
    "    x = list(range(len(df_group[f\"id_{plot_mode}\"])))\n",
    "    y = df_group['stream_p'].tolist()\n",
    "    plt.fill_between(x, y, alpha=0.6, label=r\"$P_{\\text{label}}$\", color='gray')\n",
    "\n",
    "    # Find first zero\n",
    "    zero_idx = 0\n",
    "    while y[zero_idx] > 0:\n",
    "        zero_idx += 1\n",
    "\n",
    "    if ZOOM:\n",
    "        plt.xlim(xlim[0], zero_idx)\n",
    "    else:\n",
    "        plt.axvline(zero_idx, color='gray', linestyle='--', linewidth=0.8)\n",
    "        plt.xlim(*xlim)\n",
    "    plt.axhline(0, color='gray', linestyle='-', linewidth=0.8)\n",
    "\n",
    "    labels.append(rf\"stream\")\n",
    "\n",
    "    # Plot deltas for weights: Replay vs SGD\n",
    "    colors = {\n",
    "        'sgd': sns.color_palette(\"Spectral\", 10)[0],\n",
    "        'replay': 'black',\n",
    "        'sgd_classifier_only': sns.color_palette(\"Spectral\", 10)[9],\n",
    "    }\n",
    "    linestyles = {\n",
    "        'sgd': '-',\n",
    "        'replay': (0, (1, 1)),\n",
    "        'sgd_classifier_only': '-',\n",
    "    }\n",
    "    alphas = {\n",
    "        'sgd': 1,\n",
    "        'replay': 0.6,\n",
    "        'sgd_classifier_only': 0.7,\n",
    "    }\n",
    "    if plot_weights:\n",
    "        ax.set(xlabel=plot_mode, ylabel=rf\"$\\text{{P}}(\\lVert w_{{|S|}} \\rVert - \\lVert w_{{0}} \\rVert )$\")\n",
    "        y_name = 'weight_norm_delta_p'\n",
    "    else:\n",
    "        ax.set(xlabel=plot_mode, ylabel=rf\"$\\text{{P}}(|b_{{|S|}}| - |b_{{0}}|)$\")\n",
    "        y_name = 'bias_delta_p'\n",
    "\n",
    "    for model_key in PLOT_SELECTION:\n",
    "        name = model_key\n",
    "        model_result = model_results[model_key]\n",
    "\n",
    "        df = model_result['df']\n",
    "        df_group = df.groupby([f\"id_{plot_mode}\"], as_index=False).agg({y_name: ['mean', scipy.stats.sem]})\n",
    "        df_reorder = df_group.loc[sort_idxs]\n",
    "\n",
    "        x = list(range(len(df_group[f\"id_{plot_mode}\"])))\n",
    "        y = df_reorder[y_name]['mean']\n",
    "        y_err = df_reorder[y_name]['sem']\n",
    "        label = legend_name_mapping[name]\n",
    "        plt.plot(x, y, alpha=alphas[name], linestyle=linestyles[name], color=colors[name], linewidth=1.2,\n",
    "                 label=label, )  # Takes avg over users\n",
    "\n",
    "        y_low = y - y_err\n",
    "        y_high = y + y_err\n",
    "        plt.fill_between(x, y_low, y_high, alpha=0.2, color=colors[name], linewidth=0.0)\n",
    "\n",
    "    ax.spines.right.set_visible(False)\n",
    "    ax.spines.top.set_visible(False)\n",
    "\n",
    "    leg = ax.legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    path = os.path.join(parent_dirpath, filename)\n",
    "    fig.savefig(path, bbox_inches='tight')\n",
    "\n",
    "    # PLOT BOTH\n",
    "    plt.show()\n",
    "    plt.close('all')"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}