{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interceptor\n",
    "\n",
    "In this example, you gonna learn how to use `Interceptor` to capture intermediate values in the execution of a PyTorch module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch # PyTorch\n",
    "\n",
    "from pytorch_probing import Interceptor # Intercepts intermediate values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We gonna start creating a example a module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExampleModel(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self, input_size, hidden_size, output_size, n_hidden=0):\n",
    "        super().__init__()\n",
    "        self.linear1 = torch.nn.Linear(input_size, hidden_size)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "\n",
    "        if n_hidden > 0:\n",
    "            layers = []\n",
    "            for _ in range(n_hidden):\n",
    "                layers.append(torch.nn.Linear(hidden_size, hidden_size))\n",
    "                layers.append(torch.nn.ReLU())\n",
    "            self.hidden_layers = torch.nn.Sequential(*layers)\n",
    "        self._n_hidden = n_hidden\n",
    "\n",
    "        self.linear2 = torch.nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear1(x)\n",
    "        x = self.relu(x)\n",
    "        if self._n_hidden > 0:\n",
    "            x = self.hidden_layers(x)\n",
    "        x = self.linear2(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, we gonna use 2 hidden layers between the first and last layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ExampleModel(\n",
       "  (linear1): Linear(in_features=2, out_features=3, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (hidden_layers): Sequential(\n",
       "    (0): Linear(in_features=3, out_features=3, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=3, out_features=3, bias=True)\n",
       "    (3): ReLU()\n",
       "  )\n",
       "  (linear2): Linear(in_features=3, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_size = 2\n",
    "hidden_size = 3\n",
    "output_size = 1\n",
    "n_hidden = 2\n",
    "\n",
    "model = ExampleModel(input_size, hidden_size, output_size, n_hidden)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Than, we gonna pass our created module to a Interceptor, with the paths of the submodules we wanna get its outputs. Notices that we can use \".\" to get inner submodules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "Create Interceptor"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Interceptor(\n",
       "  (_module): ExampleModel(\n",
       "    (linear1): InterceptorLayer(\n",
       "      (_module): Linear(in_features=2, out_features=3, bias=True)\n",
       "    )\n",
       "    (relu): ReLU()\n",
       "    (hidden_layers): Sequential(\n",
       "      (0): Linear(in_features=3, out_features=3, bias=True)\n",
       "      (1): ReLU()\n",
       "      (2): InterceptorLayer(\n",
       "        (_module): Linear(in_features=3, out_features=3, bias=True)\n",
       "      )\n",
       "      (3): ReLU()\n",
       "    )\n",
       "    (linear2): InterceptorLayer(\n",
       "      (_module): Linear(in_features=3, out_features=1, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paths = [\"linear1\", \"hidden_layers.2\", \"linear2\"]\n",
    "\n",
    "intercepted_model = Interceptor(model, paths, detach=False)\n",
    "intercepted_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Interceptor modifies in-place the original module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "View model"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ExampleModel(\n",
       "  (linear1): InterceptorLayer(\n",
       "    (_module): Linear(in_features=2, out_features=3, bias=True)\n",
       "  )\n",
       "  (relu): ReLU()\n",
       "  (hidden_layers): Sequential(\n",
       "    (0): Linear(in_features=3, out_features=3, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): InterceptorLayer(\n",
       "      (_module): Linear(in_features=3, out_features=3, bias=True)\n",
       "    )\n",
       "    (3): ReLU()\n",
       "  )\n",
       "  (linear2): InterceptorLayer(\n",
       "    (_module): Linear(in_features=3, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We pass a example input throught the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = torch.randn([10, 2])\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = intercepted_model(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029],\n",
       "        [-0.0029]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the interceptor captures the required outputs and stores then in the \"outputs\" attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "View intercepted outputs"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'linear1': tensor([[ 0.1654,  0.0663, -0.7493],\n",
       "         [ 0.3264, -0.9112, -0.8976],\n",
       "         [ 0.1156,  0.0056, -0.8358],\n",
       "         [ 0.7862, -1.3303, -0.4560],\n",
       "         [ 0.3589, -1.0462, -0.9049],\n",
       "         [ 0.0942, -0.3027, -0.9759],\n",
       "         [ 0.7300, -1.1132, -0.4494],\n",
       "         [ 0.5927, -0.4919, -0.4003],\n",
       "         [ 0.3150, -0.3468, -0.7065],\n",
       "         [-0.5145,  0.6613, -1.4115]]),\n",
       " 'hidden_layers.2': tensor([[-0.5792, -0.1493, -0.1676],\n",
       "         [-0.5853, -0.1682, -0.1807],\n",
       "         [-0.5907, -0.1497, -0.1603],\n",
       "         [-0.6053, -0.2082, -0.2035],\n",
       "         [-0.5843, -0.1710, -0.1838],\n",
       "         [-0.5923, -0.1483, -0.1580],\n",
       "         [-0.6003, -0.2033, -0.2024],\n",
       "         [-0.5879, -0.1913, -0.1997],\n",
       "         [-0.5856, -0.1673, -0.1795],\n",
       "         [-0.4808, -0.0937, -0.1802]]),\n",
       " 'linear2': tensor([[-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029],\n",
       "         [-0.0029]])}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intercepted_model.outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can clean this outputs with the `interceptor_clear` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": [
     "interceptor_clear example"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'linear1': None, 'hidden_layers.2': None, 'linear2': None}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intercepted_model.interceptor_clear()\n",
    "intercepted_model.outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And return the model to its original state with the reduce method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "Reduce example"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ExampleModel(\n",
       "  (linear1): Linear(in_features=2, out_features=3, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (hidden_layers): Sequential(\n",
       "    (0): Linear(in_features=3, out_features=3, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=3, out_features=3, bias=True)\n",
       "    (3): ReLU()\n",
       "  )\n",
       "  (linear2): Linear(in_features=3, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intercepted_model.reduce()\n",
    "\n",
    "model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}