d2i_modules_builders_guide.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNUIBezHodRmsPogVWZASNd",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/andrews/pytorch_tutorials/blob/main/d2i_modules_builders_guide.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dive into Deep Learning\n",
        "d2l.ai"
      ],
      "metadata": {
        "id": "1GpmAkiXkf-w"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installation"
      ],
      "metadata": {
        "id": "ui8uyYoukpT-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "torch.__version__\n",
        "# torch.cuda.current_device()\n",
        "# torch.cuda.get_device_name(0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        },
        "id": "eobdMgATkqxS",
        "outputId": "c8bd3218-0f3e-4f7e-d485-e49f36ddf37a"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'2.3.0+cu121'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.nn import functional as F"
      ],
      "metadata": {
        "id": "leSH-L5rUOi1"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))\n",
        "\n",
        "X = torch.rand(2, 20)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tCxkPYj1UPxU",
        "outputId": "cc613a5b-5f0a-4b41-d3a3-826343933e59"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
            "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n = net(X)"
      ],
      "metadata": {
        "id": "kYExmxykUSsx"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Lhn8kNNxUbjy",
        "outputId": "02f7d157-1257-416a-8c3e-2c992de18399"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class MLP(nn.Module):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.hidden = nn.LazyLinear(256)\n",
        "    self.out = nn.LazyLinear(10)\n",
        "\n",
        "  def forward(self, X):\n",
        "    return self.out(F.relu(self.hidden(X)))"
      ],
      "metadata": {
        "id": "0Z_-XFC4T_RL"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# MLP is now a module, and can be used as a layer or component in other networks/modules, and even subclassed itself.\n",
        "m = MLP()\n",
        "m(X).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S_IfyWfFZm4E",
        "outputId": "0309638d-5653-43d8-be49-f3ae975fbd7d"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
            "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# under the hood of how Sequential class actually works. Daisy-chains the modules together\n",
        "class MySequential(nn.Module):\n",
        "  def __init__(self, *args):\n",
        "    super().__init__()\n",
        "    for idx, module in enumerate(args):\n",
        "      self.add_module(str(idx), module)\n",
        "\n",
        "  def forward(self, X):\n",
        "    for module in self.children():\n",
        "      X = module(X)\n",
        "    return X"
      ],
      "metadata": {
        "id": "sGurkUVgZsc1"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ms = MySequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))\n",
        "ms(X).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U-K_ub7hboko",
        "outputId": "8583c0f5-e689-4d71-b530-70a7dade7e84"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
            "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# random constant added\n",
        "class FixedMLP(nn.Module):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.rand_weight = torch.rand((20, 20))\n",
        "    self.linear = nn.LazyLinear(20)\n",
        "\n",
        "  def forward(self, X):\n",
        "    X = self.linear(X)\n",
        "    X = F.relu(X @ self.rand_weight + 1)\n",
        "    # reuse. pass it through the linear layer again\n",
        "    # as a result the model's weights are also used twice\n",
        "    X = self.linear(X)\n",
        "    # control flow\n",
        "    while X.abs().sum() > 1:\n",
        "      X /= 2\n",
        "    return X.sum() # a scalar?? I suppose if it's the last layer"
      ],
      "metadata": {
        "id": "qbXP6Li5ecB3"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fmlp = FixedMLP()\n",
        "fmlp(X)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yBxV8GvdfVoU",
        "outputId": "4cbdfc6c-7f35-4143-adf8-920de34764ff"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
            "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(0.0759, grad_fn=<SumBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# shapeless, scalar\n",
        "# if it was [x, 1], similar to previous explanation, x would be sample size\n",
        "# 1 would be the output size of each sample\n",
        "fmlp(X).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iuDEHM8CfbY3",
        "outputId": "51a15c93-5337-4d0b-d1ae-0d268bb653dc"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([])"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# all still linear\n",
        "class NestMLP(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(nn.LazyLinear(64), nn.ReLU(),\n",
        "                                 nn.LazyLinear(32), nn.ReLU())\n",
        "        self.linear = nn.LazyLinear(16)\n",
        "\n",
        "    def forward(self, X):\n",
        "        return self.linear(self.net(X))\n",
        "\n",
        "chimera = nn.Sequential(NestMLP(), nn.LazyLinear(20), FixedMLP())\n",
        "chimera(X)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3B6d_g16gyHf",
        "outputId": "6c869eeb-4ed0-4afd-eaf2-ca12c4d0812b"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:181: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.\n",
            "  warnings.warn('Lazy modules are a new feature under heavy development '\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(0.0502, grad_fn=<SumBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# chapter question\n",
        "class ParallelMLP(nn.Module):\n",
        "  def __init__(self, net1, net2):\n",
        "    super().__init__()\n",
        "    self.hnet1 = net1\n",
        "    self.hnet2 = net2\n",
        "    self.linear = nn.LazyLinear(20, 10)\n",
        "\n",
        "  def forward(self, X1, X2):\n",
        "    h1o = self.hnet1(X1)\n",
        "    h2o = self.hnet2(X2)\n",
        "    return self.linear(torch.cat((h1o, h2o), dim=1))"
      ],
      "metadata": {
        "id": "-tyY90afhSbb"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pm = ParallelMLP(MLP(), MLP())\n",
        "pm(X, X)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MUXeUrdGkzvi",
        "outputId": "f8072ef0-0260-47e4-8c37-6d5c422c4f3a"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 0.0847, -0.0235,  0.2173, -0.1349, -0.0194,  0.1340,  0.0129, -0.1180,\n",
              "          0.1877, -0.2754, -0.2155, -0.0283,  0.0038, -0.1234,  0.3625, -0.2176,\n",
              "          0.1989,  0.1126, -0.2849,  0.0579],\n",
              "        [ 0.0672, -0.0375,  0.2277, -0.0960,  0.0018,  0.1173, -0.0297, -0.1408,\n",
              "          0.1618, -0.3578, -0.1740, -0.0145, -0.0349, -0.0846,  0.3938, -0.2156,\n",
              "          0.1601,  0.1826, -0.3229, -0.0162]], grad_fn=<AddmmBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    }
  ]
}