{
  "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: Control Flow + Weight Sharing\n======================================\n\nTo showcase the power of PyTorch dynamic graphs, we will implement a\nvery strange model: a third-fifth order polynomial that on each forward\npass chooses a random number between 4 and 5 and uses that many orders,\nreusing the same weights multiple times to compute the fourth and fifth\norder.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import random\nimport torch\nimport math\n\n\nclass DynamicNet(torch.nn.Module):\n    def __init__(self):\n        \"\"\"\n        In the constructor we instantiate five parameters and assign them as members.\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        self.e = torch.nn.Parameter(torch.randn(()))\n\n    def forward(self, x):\n        \"\"\"\n        For the forward pass of the model, we randomly choose either 4, 5\n        and reuse the e parameter to compute the contribution of these orders.\n\n        Since each forward pass builds a dynamic computation graph, we can use normal\n        Python control-flow operators like loops or conditional statements when\n        defining the forward pass of the model.\n\n        Here we also see that it is perfectly safe to reuse the same parameter many\n        times when defining a computational graph.\n        \"\"\"\n        y = self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3\n        for exp in range(4, random.randint(4, 6)):\n            y = y + self.e * x ** exp\n        return y\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 + {self.e.item()} x^4 ? + {self.e.item()} x^5 ?'\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 = DynamicNet()\n\n# Construct our loss function and an Optimizer. Training this strange model with\n# vanilla stochastic gradient descent is tough, so we use momentum\ncriterion = torch.nn.MSELoss(reduction='sum')\noptimizer = torch.optim.SGD(model.parameters(), lr=1e-8, momentum=0.9)\nfor t in range(30000):\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 % 2000 == 1999:\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
}