{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5e5f706d",
   "metadata": {
    "origin_pos": 1
   },
   "source": [
    "# Custom Layers\n",
    "\n",
    "One factor behind deep learning's success\n",
    "is the availability of a wide range of layers\n",
    "that can be composed in creative ways\n",
    "to design architectures suitable\n",
    "for a wide variety of tasks.\n",
    "For instance, researchers have invented layers\n",
    "specifically for handling images, text,\n",
    "looping over sequential data,\n",
    "and\n",
    "performing dynamic programming.\n",
    "Sooner or later, you will need\n",
    "a layer that does not exist yet in the deep learning framework.\n",
    "In these cases, you must build a custom layer.\n",
    "In this section, we show you how.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b2079e0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:10.693752Z",
     "iopub.status.busy": "2023-08-18T19:31:10.693415Z",
     "iopub.status.idle": "2023-08-18T19:31:13.986742Z",
     "shell.execute_reply": "2023-08-18T19:31:13.985398Z"
    },
    "origin_pos": 3,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from d2l import torch as d2l"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a6181d4",
   "metadata": {
    "origin_pos": 6
   },
   "source": [
    "## (**Layers without Parameters**)\n",
    "\n",
    "To start, we construct a custom layer\n",
    "that does not have any parameters of its own.\n",
    "This should look familiar if you recall our\n",
    "introduction to modules in :numref:`sec_model_construction`.\n",
    "The following `CenteredLayer` class simply\n",
    "subtracts the mean from its input.\n",
    "To build it, we simply need to inherit\n",
    "from the base layer class and implement the forward propagation function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21e03034",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:13.993899Z",
     "iopub.status.busy": "2023-08-18T19:31:13.993311Z",
     "iopub.status.idle": "2023-08-18T19:31:14.001595Z",
     "shell.execute_reply": "2023-08-18T19:31:14.000547Z"
    },
    "origin_pos": 8,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "class CenteredLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X - X.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b36f65b",
   "metadata": {
    "origin_pos": 11
   },
   "source": [
    "Let's verify that our layer works as intended by feeding some data through it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "34c473c6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.006461Z",
     "iopub.status.busy": "2023-08-18T19:31:14.005870Z",
     "iopub.status.idle": "2023-08-18T19:31:14.035296Z",
     "shell.execute_reply": "2023-08-18T19:31:14.034301Z"
    },
    "origin_pos": 12,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2., -1.,  0.,  1.,  2.])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "layer = CenteredLayer()\n",
    "layer(torch.tensor([1.0, 2, 3, 4, 5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5c4de56",
   "metadata": {
    "origin_pos": 13
   },
   "source": [
    "We can now [**incorporate our layer as a component\n",
    "in constructing more complex models.**]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da630405",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.040444Z",
     "iopub.status.busy": "2023-08-18T19:31:14.040108Z",
     "iopub.status.idle": "2023-08-18T19:31:14.044820Z",
     "shell.execute_reply": "2023-08-18T19:31:14.043922Z"
    },
    "origin_pos": 15,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "net = nn.Sequential(nn.LazyLinear(128), CenteredLayer())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a31b206",
   "metadata": {
    "origin_pos": 18
   },
   "source": [
    "As an extra sanity check, we can send random data\n",
    "through the network and check that the mean is in fact 0.\n",
    "Because we are dealing with floating point numbers,\n",
    "we may still see a very small nonzero number\n",
    "due to quantization.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "370e0abb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.048310Z",
     "iopub.status.busy": "2023-08-18T19:31:14.047972Z",
     "iopub.status.idle": "2023-08-18T19:31:14.059041Z",
     "shell.execute_reply": "2023-08-18T19:31:14.057938Z"
    },
    "origin_pos": 20,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-6.5193e-09, grad_fn=<MeanBackward0>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y = net(torch.rand(4, 8))\n",
    "Y.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44ea9d42",
   "metadata": {
    "origin_pos": 23
   },
   "source": [
    "## [**Layers with Parameters**]\n",
    "\n",
    "Now that we know how to define simple layers,\n",
    "let's move on to defining layers with parameters\n",
    "that can be adjusted through training.\n",
    "We can use built-in functions to create parameters, which\n",
    "provide some basic housekeeping functionality.\n",
    "In particular, they govern access, initialization,\n",
    "sharing, saving, and loading model parameters.\n",
    "This way, among other benefits, we will not need to write\n",
    "custom serialization routines for every custom layer.\n",
    "\n",
    "Now let's implement our own version of the  fully connected layer.\n",
    "Recall that this layer requires two parameters,\n",
    "one to represent the weight and the other for the bias.\n",
    "In this implementation, we bake in the ReLU activation as a default.\n",
    "This layer requires two input arguments: `in_units` and `units`, which\n",
    "denote the number of inputs and outputs, respectively.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d07e84dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.063246Z",
     "iopub.status.busy": "2023-08-18T19:31:14.062888Z",
     "iopub.status.idle": "2023-08-18T19:31:14.069269Z",
     "shell.execute_reply": "2023-08-18T19:31:14.068283Z"
    },
    "origin_pos": 25,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "class MyLinear(nn.Module):\n",
    "    def __init__(self, in_units, units):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.randn(in_units, units))\n",
    "        self.bias = nn.Parameter(torch.randn(units,))\n",
    "\n",
    "    def forward(self, X):\n",
    "        linear = torch.matmul(X, self.weight.data) + self.bias.data\n",
    "        return F.relu(linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "218be348",
   "metadata": {
    "origin_pos": 29,
    "tab": [
     "pytorch"
    ]
   },
   "source": [
    "Next, we instantiate the `MyLinear` class\n",
    "and access its model parameters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8a664799",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.074206Z",
     "iopub.status.busy": "2023-08-18T19:31:14.073211Z",
     "iopub.status.idle": "2023-08-18T19:31:14.080883Z",
     "shell.execute_reply": "2023-08-18T19:31:14.079861Z"
    },
    "origin_pos": 31,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.4783,  0.4284, -0.0899],\n",
       "        [-0.6347,  0.2913, -0.0822],\n",
       "        [-0.4325, -0.1645, -0.3274],\n",
       "        [ 1.1898,  0.6482, -1.2384],\n",
       "        [-0.1479,  0.0264, -0.9597]], requires_grad=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linear = MyLinear(5, 3)\n",
    "linear.weight"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7b12b77",
   "metadata": {
    "origin_pos": 34
   },
   "source": [
    "We can [**directly carry out forward propagation calculations using custom layers.**]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "859b12e2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.084676Z",
     "iopub.status.busy": "2023-08-18T19:31:14.084328Z",
     "iopub.status.idle": "2023-08-18T19:31:14.091968Z",
     "shell.execute_reply": "2023-08-18T19:31:14.090864Z"
    },
    "origin_pos": 36,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.9316, 0.0000],\n",
       "        [0.1808, 1.4208, 0.0000]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linear(torch.rand(2, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6811c960",
   "metadata": {
    "origin_pos": 39
   },
   "source": [
    "We can also (**construct models using custom layers.**)\n",
    "Once we have that we can use it just like the built-in fully connected layer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "53f3a28a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:31:14.096505Z",
     "iopub.status.busy": "2023-08-18T19:31:14.095515Z",
     "iopub.status.idle": "2023-08-18T19:31:14.104253Z",
     "shell.execute_reply": "2023-08-18T19:31:14.102782Z"
    },
    "origin_pos": 41,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000],\n",
       "        [13.0800]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))\n",
    "net(torch.rand(2, 64))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de2a529d",
   "metadata": {
    "origin_pos": 44
   },
   "source": [
    "## Summary\n",
    "\n",
    "We can design custom layers via the basic layer class. This allows us to define flexible new layers that behave differently from any existing layers in the library.\n",
    "Once defined, custom layers can be invoked in arbitrary contexts and architectures.\n",
    "Layers can have local parameters, which can be created through built-in functions.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Design a layer that takes an input and computes a tensor reduction,\n",
    "   i.e., it returns $y_k = \\sum_{i, j} W_{ijk} x_i x_j$.\n",
    "1. Design a layer that returns the leading half of the Fourier coefficients of the data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7288fd89",
   "metadata": {
    "origin_pos": 46,
    "tab": [
     "pytorch"
    ]
   },
   "source": [
    "[Discussions](https://discuss.d2l.ai/t/59)\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "required_libs": []
 },
 "nbformat": 4,
 "nbformat_minor": 5
}