register_buffer_vs_register_parameter.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNWsbU1zUPYAX3Mz3mwK1jX",
      "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/zzqnot996/learn_deeplearning/blob/main/register_buffer_vs_register_parameter.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5ocUvCMQfxbo",
        "outputId": "031758c3-9a1a-4828-d7da-ea96b939b42c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n",
            "\n",
            " ######################################## \n",
            "\n",
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class MyModule(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(MyModule,self).__init__()\n",
        "    self.conv1 = nn.Conv2d(in_channels = 3, out_channels=6, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "    self.conv2 = nn.Conv2d(in_channels = 6, out_channels=9, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "    self.weight = torch.ones(10,10)\n",
        "    self.bias = torch.zeros(10)\n",
        "\n",
        "  def forward(self,x):\n",
        "    x=self.conv1(x)\n",
        "    x=self.conv2(x)\n",
        "    x = x*self.weight + self.bias\n",
        "    return x\n",
        "\n",
        "\n",
        "net = MyModule()\n",
        "\n",
        "for name, param in net.named_parameters():\n",
        "\n",
        "  print(name,param.shape)\n",
        "\n",
        "print(\"\\n\",\"#\"*40,\"\\n\")\n",
        "\n",
        "for key, val in net.state_dict().items():\n",
        "  print(key,val.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1、`register_parameter()`\n",
        "\n",
        "`register_parameter()`是 torch.nn.Module 类中的一个方法\n",
        "\n",
        "# 1.1、 作用\n",
        "\n",
        "- 可将 `self.weight` 和 `self.bias` 定义为可学习的参数，保存到网络对象的参数中，被优化器作用进行学习\n",
        "    \n",
        "- `self.weight` 和 `self.bias` 可被保存到 state\\_dict 中，进而可以 保存到网络文件 / 网络参数文件中\n",
        "    \n",
        "\n",
        "# 1.2、用法\n",
        "\n",
        "  `register_parameter(name，param)`\n",
        "\n",
        "- `name`：参数名称\n",
        "    \n",
        "- `param`：参数张量， 须是 `torch.nn.Parameter()` 对象 或 None ，否则报错如下\n",
        "    \n",
        "\n",
        "> TypeError: cannot assign 'torch.FloatTensor' object to parameter 'xx' (torch.nn.Parameter or None required)"
      ],
      "metadata": {
        "id": "Y_lQwiIckx7v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class MyModule(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyModule, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.conv2 = nn.Conv2d(in_channels=6, out_channels=9, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "\n",
        "        # register_parameter  -----must torch.nn.Parameter\n",
        "        self.register_parameter('weight', torch.nn.Parameter(torch.ones(10, 10)))\n",
        "        self.register_parameter('bias', torch.nn.Parameter(torch.zeros(10)))\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "        x = self.conv2(x)\n",
        "        x = x * self.weight + self.bias\n",
        "        return x\n",
        "\n",
        "\n",
        "net = MyModule()\n",
        "\n",
        "for name, param in net.named_parameters():\n",
        "    print(name, param.shape)\n",
        "\n",
        "print('\\n', '*'*40, '\\n')\n",
        "\n",
        "for key, val in net.state_dict().items():\n",
        "    print(key, val.shape)\n",
        ""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Kl2Aw1G9k0hJ",
        "outputId": "ced8f9ef-fe1d-4ba7-ef99-43df66d01fec"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "weight torch.Size([10, 10])\n",
            "bias torch.Size([10])\n",
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n",
            "\n",
            " **************************************** \n",
            "\n",
            "weight torch.Size([10, 10])\n",
            "bias torch.Size([10])\n",
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2、`register_buffer()`\n",
        "\n",
        "`register_buffer()`是 `torch.nn.Module()` 类中的一个方法\n",
        "\n",
        "# 2.1 、作用\n",
        "\n",
        "- 将 `self.weight` 和 `self.bias` 定义为不可学习的参数，不会被保存到网络对象的参数中，不会被优化器作用进行学习\n",
        "    \n",
        "- `self.weight` 和 `self.bias` 可被保存到 state\\_dict 中，进而可以 保存到网络文件 / 网络参数文件中\n",
        "    \n",
        "\n",
        "它用于在网络实例中 注册缓冲区，存储在缓冲区中的数据，类似于参数（但不是参数）\n",
        "\n",
        "- 参数：可以被优化器更新  （requires\\_grad=False / True）\n",
        "    \n",
        "- buffer 中的数据 ： 不会被优化器更新\n",
        "    \n",
        "\n",
        "# 2.2、用法\n",
        "\n",
        "  `register_buffer(name，tensor)`\n",
        "\n",
        "- `name`：参数名称\n",
        "    \n",
        "- `tensor`：张量"
      ],
      "metadata": {
        "id": "MMwl31RwlFCV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class MyModule(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyModule, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.conv2 = nn.Conv2d(in_channels=6, out_channels=9, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "\n",
        "        self.register_buffer('weight', torch.ones(10, 10))\n",
        "        self.register_buffer('bias', torch.zeros(10))\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "        x = self.conv2(x)\n",
        "        x = x * self.weight + self.bias\n",
        "        return x\n",
        "\n",
        "\n",
        "net = MyModule()\n",
        "\n",
        "for name, param in net.named_parameters():\n",
        "    print(name, param.shape)\n",
        "\n",
        "print('\\n', '*'*40, '\\n')\n",
        "\n",
        "for key, val in net.state_dict().items():\n",
        "    print(key, val.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pj1oItdflGAj",
        "outputId": "e9e2182b-2502-47c2-cb37-2d9d15dafccb"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n",
            "\n",
            " **************************************** \n",
            "\n",
            "weight torch.Size([10, 10])\n",
            "bias torch.Size([10])\n",
            "conv1.weight torch.Size([6, 3, 3, 3])\n",
            "conv2.weight torch.Size([9, 6, 3, 3])\n"
          ]
        }
      ]
    }
  ]
}