AddMul-Experiment1-InspectMasks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5517b936",
   "metadata": {},
   "source": [
    "# Uncovering Masked Network for particular subtasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0788e6f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from models.addmul import HandleAddMul"
   ]
  },
  {
   "cell_type": "raw",
   "id": "96b29cec",
   "metadata": {},
   "source": [
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37a4f0a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "network_cache_dir = \"networks/cache-networks/\"\n",
    "network_name = \"lyr256-split0.8-lr0.01-add-mul.data\"\n",
    "\n",
    "checkpoint = True\n",
    "test_flag = 1\n",
    "\n",
    "input_dims = [42]\n",
    "output_dims = [20]\n",
    "batchsize = 128\n",
    "tau = 1  # temperature parameter, NB: check for value in paper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a2ec95c",
   "metadata": {},
   "source": [
    "# Test the Add mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "225b7427",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "... FNN Network training on cuda:0 ...\n",
      "Accessing : networks/cache-networks/lyr256-split0.8-lr0.01-add-mul.data\n",
      "networks/cache-networks/lyr256-split0.8-lr0.01-add-mul.data\n"
     ]
    }
   ],
   "source": [
    "handler = HandleAddMul(input_dims, output_dims, dir=network_cache_dir + network_name, checkpoint=checkpoint, lr=0.001)\n",
    "handler.refreeze_weights()\n",
    "\n",
    "# Fetch Mask\n",
    "load_logits = torch.load('masks/trained_logits_add_mask_.pt')\n",
    "\n",
    "# Convert logits into binary mask\n",
    "binary_mask = []\n",
    "all_logits = []\n",
    "for layer in load_logits:\n",
    "    sig = torch.sigmoid(layer)\n",
    "    greaterthan = torch.ge(sig, 0.5) # Bool mask\n",
    "    all_logits.append(torch.masked_select(sig, greaterthan).view(-1).cpu().detach().numpy())\n",
    "    binary_mask.append(greaterthan.float()) # Binary mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dbaf9e5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_split = 0.8\n",
    "test_split = 1 - train_split\n",
    "\n",
    "data_fp = [\"generate_datasets/tmp/digit-data/simple_add.npy\",\n",
    "           \"generate_datasets/tmp/digit-data/simple_mul.npy\"]\n",
    "data_add = np.load(data_fp[0], allow_pickle=True)\n",
    "data_mul = np.load(data_fp[1], allow_pickle=True)\n",
    "data_add_len = len(data_add)\n",
    "data_mul_len = len(data_mul)\n",
    "\n",
    "train_split_idx_add = int(data_add_len * train_split)\n",
    "train_split_idx_mul = int(data_mul_len * train_split)\n",
    "test_data_add = data_add[train_split_idx_add:]\n",
    "test_data_mul = data_mul[train_split_idx_mul:]\n",
    "\n",
    "test_loader_add = torch.utils.data.DataLoader(dataset=torch.Tensor(test_data_add), batch_size=batchsize, shuffle=True)\n",
    "test_loader_mul = torch.utils.data.DataLoader(dataset=torch.Tensor(test_data_mul), batch_size=batchsize, shuffle=True)\n",
    "\n",
    "def cycle(iterable):\n",
    "    while True:\n",
    "        for x in iterable:\n",
    "            yield x\n",
    "\n",
    "iterator_test_add = iter(cycle(test_loader_add))\n",
    "iterator_test_mul = iter(cycle(test_loader_mul))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "93f25cc9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting epoch 0...\n",
      "Starting epoch 1...\n",
      "Starting epoch 2...\n",
      "Starting epoch 3...\n",
      "Starting epoch 4...\n",
      "Starting epoch 5...\n",
      "Starting epoch 6...\n",
      "Starting epoch 7...\n",
      "Starting epoch 8...\n",
      "Starting epoch 9...\n",
      "Starting epoch 10...\n",
      "Starting epoch 11...\n",
      "Starting epoch 12...\n",
      "Starting epoch 13...\n",
      "Starting epoch 14...\n",
      "Starting epoch 15...\n",
      "Starting epoch 16...\n",
      "Starting epoch 17...\n",
      "Starting epoch 18...\n",
      "Starting epoch 19...\n",
      "Starting epoch 20...\n",
      "Starting epoch 21...\n",
      "Starting epoch 22...\n",
      "Starting epoch 23...\n",
      "Starting epoch 24...\n",
      "Starting epoch 25...\n",
      "Starting epoch 26...\n",
      "Starting epoch 27...\n",
      "Starting epoch 28...\n",
      "Starting epoch 29...\n",
      "Starting epoch 30...\n",
      "Starting epoch 31...\n",
      "Starting epoch 32...\n",
      "Starting epoch 33...\n",
      "Starting epoch 34...\n",
      "Starting epoch 35...\n",
      "Starting epoch 36...\n",
      "Starting epoch 37...\n",
      "Starting epoch 38...\n",
      "Starting epoch 39...\n",
      "Starting epoch 40...\n",
      "Starting epoch 41...\n",
      "Starting epoch 42...\n",
      "Starting epoch 43...\n",
      "Starting epoch 44...\n",
      "Starting epoch 45...\n",
      "Starting epoch 46...\n",
      "Starting epoch 47...\n",
      "Starting epoch 48...\n",
      "Starting epoch 49...\n",
      "Starting epoch 50...\n",
      "Starting epoch 51...\n",
      "Starting epoch 52...\n",
      "Starting epoch 53...\n",
      "Starting epoch 54...\n",
      "Starting epoch 55...\n",
      "Starting epoch 56...\n",
      "Starting epoch 57...\n",
      "Starting epoch 58...\n",
      "Starting epoch 59...\n",
      "Starting epoch 60...\n",
      "Starting epoch 61...\n",
      "Starting epoch 62...\n",
      "Starting epoch 63...\n",
      "Starting epoch 64...\n",
      "Starting epoch 65...\n",
      "Starting epoch 66...\n",
      "Starting epoch 67...\n",
      "Starting epoch 68...\n",
      "Starting epoch 69...\n",
      "Starting epoch 70...\n",
      "Starting epoch 71...\n",
      "Starting epoch 72...\n",
      "Starting epoch 73...\n",
      "Starting epoch 74...\n",
      "Starting epoch 75...\n",
      "Starting epoch 76...\n",
      "Starting epoch 77...\n",
      "Starting epoch 78...\n",
      "Starting epoch 79...\n",
      "Starting epoch 80...\n",
      "Starting epoch 81...\n",
      "Starting epoch 82...\n",
      "Starting epoch 83...\n",
      "Starting epoch 84...\n",
      "Starting epoch 85...\n",
      "Starting epoch 86...\n",
      "Starting epoch 87...\n",
      "Starting epoch 88...\n",
      "Starting epoch 89...\n",
      "Starting epoch 90...\n",
      "Starting epoch 91...\n",
      "Starting epoch 92...\n",
      "Starting epoch 93...\n",
      "Starting epoch 94...\n",
      "Starting epoch 95...\n",
      "Starting epoch 96...\n",
      "Starting epoch 97...\n",
      "Starting epoch 98...\n",
      "Starting epoch 99...\n",
      "Add Accuracy: 0.9961239035087719 \n",
      " Mul Accuracy: 0.8756026785714285\n"
     ]
    }
   ],
   "source": [
    "loss_hist = []\n",
    "NUM_EPOCHS = 100\n",
    "\n",
    "mul_acc = 0.\n",
    "add_acc = 0.\n",
    "\n",
    "with torch.no_grad():\n",
    "    iterator_binaries = iter(binary_mask)\n",
    "    for layer in handler.network.layers[0]:\n",
    "                if isinstance(layer, torch.nn.Linear):\n",
    "                    layer.weight.mul_(next(iterator_binaries))\n",
    "\n",
    "for e in range(NUM_EPOCHS):\n",
    "    print(f'Starting epoch {e}...')\n",
    "    '''Reload weights'''\n",
    "    handler.network.load_save()\n",
    "    handler.refreeze_weights()\n",
    "    \n",
    "    '''Call Addition Mask'''\n",
    "    batch_add = next(iterator_test_add)\n",
    "    batch_mul = next(iterator_test_mul)\n",
    "    iterator_binary = iter(binary_mask)\n",
    "    \n",
    "    batches = [batch_add, batch_mul]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        res = [] \n",
    "        for batch in batches:\n",
    "            inp = torch.stack([torch.stack([b[0], b[1]]) for b in batch])\n",
    "            otp = torch.stack([b[2] for b in batch])\n",
    "            ops = torch.stack([b[3] for b in batch])\n",
    "            # Convert batch data toone-hot representation\n",
    "            inp, _ = handler.set_batched_digits(inp, otp, ops)\n",
    "\n",
    "            inp_ = inp.to(handler.network.device)\n",
    "            \n",
    "            otp = otp.to(handler.network.device)\n",
    "\n",
    "            otp_pred = handler.network(inp_)\n",
    "            \n",
    "            otp_pred = torch.stack([otp_pred[:,:10], otp_pred[:,10:]], dim=1)\n",
    "            otp_pred = torch.argmax(otp_pred, dim=2)\n",
    "            otp_pred = otp_pred[:, 0]*10 + otp_pred[:, 1]\n",
    "            res.append([otp_pred, otp])\n",
    "    \n",
    "        for i, results in enumerate(res):\n",
    "            \n",
    "            diff = results[1] - results[0]\n",
    "            running_acc = float(len(diff[diff == 0])/(diff.size()[0]))\n",
    "            if not i:\n",
    "                add_acc += running_acc\n",
    "            else:\n",
    "                mul_acc += running_acc\n",
    "\n",
    "print(f'Add Accuracy: {add_acc/NUM_EPOCHS} \\n Mul Accuracy: {mul_acc/NUM_EPOCHS}')"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}