{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fbda807b",
   "metadata": {},
   "source": [
    "# Observe model variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa0a8720",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "from dataset.lenet_cifar import LeNet, make_testloader\n",
    "from mrfi import MRFI, EasyConfig\n",
    "from mrfi.experiment import get_activation_info, get_weight_info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f999fc",
   "metadata": {},
   "source": [
    "Load MRFI model with empty config for observe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5d3625de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<mrfi.mrfi.MRFI object at 0x000001CE42E4F040> Net(\n",
      "  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (fc1): Linear(in_features=400, out_features=120, bias=True)\n",
      "  (fc2): Linear(in_features=120, out_features=84, bias=True)\n",
      "  (fc3): Linear(in_features=84, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "fi_model = MRFI(LeNet(trained=True).eval(), EasyConfig())\n",
    "print(fi_model, fi_model.model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a8ce62",
   "metadata": {},
   "source": [
    "Then make cifar10 testloader, verify input shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3f1975e0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 32, 32]) label: tensor(3)\n"
     ]
    }
   ],
   "source": [
    "input_images = make_testloader(128, batch_size = 128)\n",
    "imgs, labels = next(iter(input_images))\n",
    "print(imgs[0].shape, 'label:', labels[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e047d4e8",
   "metadata": {},
   "source": [
    "Firstly, lets observe feature map shape by simply call `get_activation_info` with `method = Shape`!\n",
    "Also, we can observe weights in same way. `get_weight_info` does not require inference.\n",
    "Observer methods are defined in mrfi.observer, you can also add custom observer by `mrfi.add_function`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "12c2a463",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation Shape:\n",
      "{'model.Shape': torch.Size([1, 10]),\n",
      " 'model.conv1.Shape': torch.Size([1, 6, 28, 28]),\n",
      " 'model.conv2.Shape': torch.Size([1, 16, 10, 10]),\n",
      " 'model.fc1.Shape': torch.Size([1, 120]),\n",
      " 'model.fc2.Shape': torch.Size([1, 84]),\n",
      " 'model.fc3.Shape': torch.Size([1, 10])}\n",
      "'Weight Shape:'\n",
      "{'conv1.weight': torch.Size([6, 3, 5, 5]),\n",
      " 'conv2.weight': torch.Size([16, 6, 5, 5]),\n",
      " 'fc1.weight': torch.Size([120, 400]),\n",
      " 'fc2.weight': torch.Size([84, 120]),\n",
      " 'fc3.weight': torch.Size([10, 84])}\n"
     ]
    }
   ],
   "source": [
    "print(\"Activation Shape:\")\n",
    "pprint(get_activation_info(fi_model, torch.zeros([1,3,32,32]), 'Shape'))\n",
    "pprint(\"Weight Shape:\")\n",
    "pprint(get_weight_info(fi_model, 'Shape'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a67ccde",
   "metadata": {},
   "source": [
    "Then, we can observe min/max ranges of values, this is very useful to make a quantization model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e2e5e515",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation MinMax:\n",
      "{'model.MinMax': (-10.844356536865234, 8.558104515075684),\n",
      " 'model.conv1.MinMax': (-5.00702428817749, 5.420466423034668),\n",
      " 'model.conv2.MinMax': (-9.770880699157715, 6.261801242828369),\n",
      " 'model.fc1.MinMax': (-13.205606460571289, 9.709344863891602),\n",
      " 'model.fc2.MinMax': (-8.901363372802734, 8.150411605834961),\n",
      " 'model.fc3.MinMax': (-10.844356536865234, 8.558104515075684)}\n",
      "Weight MinMax:\n",
      "{'conv1.weight': (-0.3178213834762573, 0.3564322888851166),\n",
      " 'conv2.weight': (-0.3756081163883209, 0.32203108072280884),\n",
      " 'fc1.weight': (-0.6651991009712219, 0.4576374590396881),\n",
      " 'fc2.weight': (-0.4559188485145569, 0.33350828289985657),\n",
      " 'fc3.weight': (-0.44256511330604553, 0.4379429221153259)}\n"
     ]
    }
   ],
   "source": [
    "print(\"Activation MinMax:\")\n",
    "pprint(get_activation_info(fi_model, input_images))\n",
    "print(\"Weight MinMax:\")\n",
    "pprint(get_weight_info(fi_model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e855ab7",
   "metadata": {},
   "source": [
    "`MaxAbs` is a similiar observer.\n",
    "We can filter modules by optional argument `module_name`, `module_type` or `module_fullname`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "508f94ce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation MaxAbs:\n",
      "{'model.conv1.MaxAbs': 5.420466423034668,\n",
      " 'model.conv2.MaxAbs': 9.770880699157715}\n",
      "Weight MaxAbs:\n",
      "{'fc1.weight': 0.6651991009712219,\n",
      " 'fc2.weight': 0.4559188485145569,\n",
      " 'fc3.weight': 0.44256511330604553}\n"
     ]
    }
   ],
   "source": [
    "print(\"Activation MaxAbs:\")\n",
    "pprint(get_activation_info(fi_model, input_images, 'MaxAbs', module_name = ['conv']))\n",
    "print(\"Weight MaxAbs:\")\n",
    "pprint(get_weight_info(fi_model, 'MaxAbs', module_type = ['Linear']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd78766",
   "metadata": {},
   "source": [
    "To obtain how divergent the distribution of variables is, let's\n",
    "observe standard deviation of variable by `Std` observer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ea062af7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation Std:\n",
      "{'model.Std': 3.0110209209484347,\n",
      " 'model.conv1.Std': 0.6536869451555304,\n",
      " 'model.conv2.Std': 0.960684580199324,\n",
      " 'model.fc1.Std': 2.5316612945593397,\n",
      " 'model.fc2.Std': 1.8289168509494267,\n",
      " 'model.fc3.Std': 3.0110209209484347}\n",
      "Weight Std:\n",
      "{'conv1.weight': 0.12276468009677341,\n",
      " 'conv2.weight': 0.10105224038704899,\n",
      " 'fc1.weight': 0.07574824678816457,\n",
      " 'fc2.weight': 0.08999676217606199,\n",
      " 'fc3.weight': 0.13330479151244343}\n"
     ]
    }
   ],
   "source": [
    "print(\"Activation Std:\")\n",
    "pprint(get_activation_info(fi_model, input_images, 'Std'))\n",
    "print(\"Weight Std:\")\n",
    "pprint(get_weight_info(fi_model, 'Std'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8142192",
   "metadata": {},
   "source": [
    "Next, let's simply draw variable distribution figure by `UniformSampling` observer!\n",
    "This method defaults to uniformly sampling up to 10000 data points, \n",
    "so you don't have to worry about memory overflow due to too large tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "16990253",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bodies': [<matplotlib.collections.PolyCollection at 0x1ce7362ef10>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73729070>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce737293d0>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73729730>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73729a90>],\n",
       " 'cmaxes': <matplotlib.collections.LineCollection at 0x1ce7362ee50>,\n",
       " 'cmins': <matplotlib.collections.LineCollection at 0x1ce7373a220>,\n",
       " 'cbars': <matplotlib.collections.LineCollection at 0x1ce7373a640>}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 51488 characters]",
      "text/plain": [
       "<Figure size 800x550 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "act_samping = get_activation_info(fi_model, input_images, 'UniformSampling', module_name = ['conv', 'fc'])\n",
    "fig, axs = plt.subplots(2,3)\n",
    "for i, (name, value) in enumerate(act_samping.items()):\n",
    "    ax = axs[i//3, i%3]\n",
    "    ax.hist(value, bins=20)\n",
    "    ax.set_title(name)\n",
    "axs[1,2].violinplot(list(act_samping.values()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8337e9d",
   "metadata": {},
   "source": [
    "Draw weights distribution by same way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1f8320d7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bodies': [<matplotlib.collections.PolyCollection at 0x1ce73857d30>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73bb2070>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73bb23d0>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73bb2730>,\n",
       "  <matplotlib.collections.PolyCollection at 0x1ce73bb2a90>],\n",
       " 'cmaxes': <matplotlib.collections.LineCollection at 0x1ce73857d60>,\n",
       " 'cmins': <matplotlib.collections.LineCollection at 0x1ce73bc0220>,\n",
       " 'cbars': <matplotlib.collections.LineCollection at 0x1ce73bc0640>}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52016 characters]",
      "text/plain": [
       "<Figure size 800x550 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "weight_samping = get_weight_info(fi_model, 'UniformSampling')\n",
    "fig, axs = plt.subplots(2,3)\n",
    "for i, (name, value) in enumerate(weight_samping.items()):\n",
    "    ax = axs[i//3, i%3]\n",
    "    ax.hist(value, bins=20)\n",
    "    ax.set_title(name)\n",
    "axs[1,2].violinplot(list(weight_samping.values()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27a058f9",
   "metadata": {},
   "source": [
    "Finally, let's visualize featuremap by simple by `SaveLast` method.\n",
    "This observer saves inference variable,\n",
    "So we expect get featuremap size (batch_size = 128, channels = 6, height = 28, width = 28) of LeNet conv1 output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dea8ad32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['model.conv1.SaveLast'])\n"
     ]
    }
   ],
   "source": [
    "last_featuremap = get_activation_info(fi_model, input_images, 'SaveLast', module_name = ['conv1'])\n",
    "print(last_featuremap.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8c4cbf8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(128, 6, 28, 28)\n"
     ]
    }
   ],
   "source": [
    "featuremap = last_featuremap['model.conv1.SaveLast'][1].numpy()\n",
    "print(featuremap.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d5df91c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1ce73c9aaf0>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 42444 characters]",
      "text/plain": [
       "<Figure size 800x550 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(featuremap[0].reshape(-1, 28))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a05817",
   "metadata": {},
   "source": [
    "Done! We have made common observations on the model, is it simple and clear enough? :)\n",
    "\n",
    "Another type of observer, including `RMSE`/`MAE`/`EqualRate` ,can compare the impact of fault injections. They needs to be used in conjunction with the fault injector."
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
