{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# For tips on running notebooks in Google Colab, see\n# https://codelin.vip/beginner/colab\n%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "PyTorch: Custom nn Modules\n==========================\n\nA third order polynomial, trained to predict $y=\\sin(x)$ from $-\\pi$ to\n$\\pi$ by minimizing squared Euclidean distance.\n\nThis implementation defines the model as a custom Module subclass.\nWhenever you want a model more complex than a simple sequence of\nexisting Modules you will need to define your model this way.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nimport math\n\n\nclass Polynomial3(torch.nn.Module):\n    def __init__(self):\n        \"\"\"\n        In the constructor we instantiate four parameters and assign them as\n        member parameters.\n        \"\"\"\n        super().__init__()\n        self.a = torch.nn.Parameter(torch.randn(()))\n        self.b = torch.nn.Parameter(torch.randn(()))\n        self.c = torch.nn.Parameter(torch.randn(()))\n        self.d = torch.nn.Parameter(torch.randn(()))\n\n    def forward(self, x):\n        \"\"\"\n        In the forward function we accept a Tensor of input data and we must return\n        a Tensor of output data. We can use Modules defined in the constructor as\n        well as arbitrary operators on Tensors.\n        \"\"\"\n        return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3\n\n    def string(self):\n        \"\"\"\n        Just like any class in Python, you can also define custom method on PyTorch modules\n        \"\"\"\n        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'\n\n\n# Create Tensors to hold input and outputs.\nx = torch.linspace(-math.pi, math.pi, 2000)\ny = torch.sin(x)\n\n# Construct our model by instantiating the class defined above\nmodel = Polynomial3()\n\n# Construct our loss function and an Optimizer. The call to model.parameters()\n# in the SGD constructor will contain the learnable parameters (defined \n# with torch.nn.Parameter) which are members of the model.\ncriterion = torch.nn.MSELoss(reduction='sum')\noptimizer = torch.optim.SGD(model.parameters(), lr=1e-6)\nfor t in range(2000):\n    # Forward pass: Compute predicted y by passing x to the model\n    y_pred = model(x)\n\n    # Compute and print loss\n    loss = criterion(y_pred, y)\n    if t % 100 == 99:\n        print(t, loss.item())\n\n    # Zero gradients, perform a backward pass, and update the weights.\n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n\nprint(f'Result: {model.string()}')"
      ]
    }
  ],
  "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}