Pytorch – 重みの初期化手法と各モジュールのデフォルトの初期化方法

1. 概要
2. torch.nn.init.uniform_ – 一様分布
3. torch.nn.init.normal_ – 正規分布
4. torch.nn.init.constant_ – 定数
5. torch.nn.init.ones_ – 定数1
6. torch.nn.init.zeros_ – 定数0
7. torch.nn.init.eye_ – 単位行列
8. torch.nn.init.dirac_ – ディラックのデルタ関数
9. torch.nn.init.xavieruniform – Xavier の方法 (一様分布)
- 9.1. fan_in、fan_out について
10. torch.nn.init.xaviernormal – Xavier の方法 (正規分布)
11. torch.nn.init.kaiminguniform – He の方法 (一様分布)
12. torch.nn.init.kaimingnormal – He の方法 (正規分布)
13. 各モジュールのデフォルトの初期化方法

概要

Pytroch に実装されているパラメータの初期化方法について解説します。

torch.nn.init.uniform_ – 一様分布

torch.nn.init.uniform_(tensor, a=0.0, b=1.0) はテンソルを $[a, b]$ の一様分布で初期化する関数です。

In [1]:

import torch

x = torch.empty(10)
torch.nn.init.uniform_(x, a=0, b=5)
print(x)

tensor([0.0789, 2.7473, 3.9215, 3.5201, 3.9573, 4.4034, 1.6722, 2.9046, 4.6709,
        2.1375])

torch.nn.init.normal_ – 正規分布

torch.nn.init.normal_(tensor, mean=0.0, std=1.0) はテンソルを平均 mean、分散 std**2 の正規分布で初期化する関数です。

In [2]:

import torch

x = torch.empty(10)
torch.nn.init.normal_(x, mean=0, std=1)
print(x)

tensor([-0.8716,  1.1498, -0.9712,  0.0973, -0.6591,  0.4944,  1.1363, -0.5743,
        -0.2399,  0.1864])

torch.nn.init.constant_ – 定数

torch.nn.init.constant_(tensor, val) はテンソルを定数 val で初期化する関数です。

In [3]:

import torch

x = torch.empty(10)
torch.nn.init.constant_(x, 1)
print(x)

tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

torch.nn.init.ones_ – 定数1

torch.nn.init.ones_(tensor) はテンソルを定数1で初期化する関数です。

In [4]:

import torch

x = torch.empty(10)
torch.nn.init.ones_(x)
print(x)

tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

torch.nn.init.zeros_ – 定数0

torch.nn.init.zeros_(tensor) はテンソルを定数0で初期化する関数です。

In [5]:

import torch

x = torch.empty(10)
torch.nn.init.zeros_(x)
print(x)

tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

torch.nn.init.eye_ – 単位行列

torch.nn.init.eye_(tensor) は2次元のテンソルを単位行列で初期化する関数です。

In [6]:

import torch

x = torch.empty(3, 3)
torch.nn.init.eye_(x)
print(x)

x = torch.empty(5, 7)
torch.nn.init.eye_(x)
print(x)

tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])
tensor([[1., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0.]])

torch.nn.init.dirac_ – ディラックのデルタ関数

torch.nn.init.dirac_(tensor) は3，4、5次元のテンソルをディラックのデルタ関数で初期化する関数です。

In [7]:

import torch

x = torch.empty(3, 3, 3)
torch.nn.init.dirac_(x)
print(x)

tensor([[[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 1., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 0., 0.],
         [0., 1., 0.]]])

torch.nn.init.xavieruniform – Xavier の方法 (一様分布)

torch.nn.init.xavier_uniform_(tensor, gain=1.0) は Xavier の方法 (Glorot 初期化ともいう) の一様分布で初期化する関数です。

$$ a = \text{gain} \times \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}} $$

としたとき、$[-a, a]$ の一様分布で初期化を行うものです。

fan_in、fan_out について

畳み込み層の重みの形状を (out_channels, in_channels, *kernel_size) としたとき、fan_in、fan_out とは次のように計算される値です。ただし、kernel_size はカーネルのサイズを表すタプルで Conv1d なら (k1,)、Conv2d なら (k1, k2)、Conv3d なら (k1, k2, k3) になります。prod は累乗 (例: prod((k1, k2, k3)) = k1 * k2 * k3) を表します。

fan_in = in_channels * prod(kernel_size)
fan_out = out_channels * prod(kernel_size)

In [8]:

import math


def calculate_fan_in_and_fan_out(tensor):
    assert tensor.dim() >= 3
    in_channels = tensor.size(1)
    out_channels = tensor.size(0)
    kernel_size = math.prod(tensor.size()[2:])

    fan_in = in_channels * kernel_size
    fan_out = out_channels * kernel_size

    return fan_in, fan_out


conv = nn.Conv1d(in_channels=5, out_channels=10, kernel_size=3)
fan_in, fan_out = calculate_fan_in_and_fan_out(conv.weight)
print(f"Conv1d shape: {conv.weight.shape}, fan_in: {fan_in}, fan_out: {fan_out}")

conv = nn.Conv2d(in_channels=5, out_channels=10, kernel_size=3)
fan_in, fan_out = calculate_fan_in_and_fan_out(conv.weight)
print(f"Conv2d shape: {conv.weight.shape}, fan_in: {fan_in}, fan_out: {fan_out}")

conv = nn.Conv3d(in_channels=5, out_channels=10, kernel_size=3)
fan_in, fan_out = calculate_fan_in_and_fan_out(conv.weight)
print(f"Conv3d shape: {conv.weight.shape}, fan_in: {fan_in}, fan_out: {fan_out}")

Conv1d shape: torch.Size([10, 5, 3]), fan_in: 15, fan_out: 30
Conv2d shape: torch.Size([10, 5, 3, 3]), fan_in: 45, fan_out: 90
Conv3d shape: torch.Size([10, 5, 3, 3, 3]), fan_in: 135, fan_out: 270

In [9]:

import torch.nn as nn

conv = nn.Conv1d(in_channels=3, out_channels=4, kernel_size=3)

torch.nn.init.xavier_uniform_(conv.weight, gain=1.0)

Parameter containing:
tensor([[[-0.2460,  0.2175, -0.5319],
         [-0.3799, -0.4568,  0.2370],
         [ 0.3525, -0.3530, -0.4080]],

        [[-0.4691, -0.1311, -0.4520],
         [ 0.1070,  0.1376, -0.3987],
         [ 0.3074,  0.2339,  0.3000]],

        [[-0.2951, -0.3920,  0.1775],
         [ 0.0544,  0.4971,  0.0906],
         [-0.5228, -0.0251,  0.3910]],

        [[-0.4156, -0.0423,  0.0640],
         [ 0.0202,  0.5138,  0.1344],
         [ 0.4507,  0.4491, -0.4352]]], requires_grad=True)

torch.nn.init.xaviernormal – Xavier の方法 (正規分布)

torch.nn.init.xavier_normal_(tensor, gain=1.0) は Xavier の方法 (Glorot 初期化ともいう) の正規分布で初期化する関数です。

$$ std = \text{gain} \times \sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}} $$

としたとき、平均0、分散 std**2 の正規分布で初期化する関数です。

In [10]:

import torch.nn as nn

conv = nn.Conv1d(in_channels=3, out_channels=4, kernel_size=3)

torch.nn.init.xavier_normal_(conv.weight, gain=1.0)

Parameter containing:
tensor([[[-0.2056,  0.0869, -0.0975],
         [ 0.0291, -0.0659,  0.1770],
         [-0.3691,  0.0374,  0.2935]],

        [[ 0.2736,  0.0647,  0.2943],
         [ 0.0235, -0.6853,  0.0501],
         [-0.0109, -0.1333, -0.1458]],

        [[ 0.1033,  0.3518,  0.0459],
         [ 0.2188,  0.3011,  0.9224],
         [ 0.0318,  0.2838,  0.2632]],

        [[-0.0945,  0.2279, -0.1813],
         [ 0.0975, -0.2411, -0.0182],
         [ 0.4383, -0.1586, -0.2009]]], requires_grad=True)

torch.nn.init.kaiminguniform – He の方法 (一様分布)

torch.nn.init.kaiming_uniform_(tensor, a=0, mode="fan_in", nonlinearity="leaky_relu") は He の方法の一様分布で初期化する関数です。

$$ a = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} $$

としたとき、$[-a, a]$ の一様分布で初期化を行うものです。

In [11]:

import torch.nn as nn

conv = nn.Conv1d(in_channels=3, out_channels=4, kernel_size=3)

torch.nn.init.kaiming_uniform_(conv.weight, mode="fan_in", nonlinearity="relu")

Parameter containing:
tensor([[[-0.5564,  0.2750, -0.4595],
         [ 0.5681,  0.4487, -0.6639],
         [ 0.1159,  0.2904,  0.4377]],

        [[-0.5020,  0.1849,  0.5438],
         [ 0.1629,  0.5750, -0.6051],
         [-0.5836, -0.0322,  0.0047]],

        [[ 0.7154,  0.2076, -0.2785],
         [-0.0636, -0.0628, -0.7857],
         [ 0.7187,  0.4702,  0.1969]],

        [[ 0.5643,  0.7418, -0.4367],
         [-0.2407, -0.3289, -0.6704],
         [-0.3830, -0.5113,  0.2795]]], requires_grad=True)

torch.nn.init.kaimingnormal – He の方法 (正規分布)

torch.nn.init.kaiming_normal_(tensor, a=0, mode="fan_in", nonlinearity="leaky_relu") は He の方法の正規分布で初期化する関数です。

$$ std = \frac{\text{gain}}{\sqrt{\text{fan\_mode}}} $$

としたとき、平均0、分散 std**2 の正規分布で初期化する関数です。

In [12]:

import torch.nn as nn

conv = nn.Conv1d(in_channels=3, out_channels=4, kernel_size=3)

torch.nn.init.kaiming_normal_(conv.weight)

Parameter containing:
tensor([[[-0.2056,  0.0869, -0.0975],
         [ 0.0291, -0.0659,  0.1770],
         [-0.3691,  0.0374,  0.2935]],

        [[ 0.2736,  0.0647,  0.2943],
         [ 0.0235, -0.6853,  0.0501],
         [-0.0109, -0.1333, -0.1458]],

        [[ 0.1033,  0.3518,  0.0459],
         [ 0.2188,  0.3011,  0.9224],
         [ 0.0318,  0.2838,  0.2632]],

        [[-0.0945,  0.2279, -0.1813],
         [ 0.0975, -0.2411, -0.0182],
         [ 0.4383, -0.1586, -0.2009]]], requires_grad=True)

各モジュールのデフォルトの初期化方法

畳み込み層

nn.Conv1d、nn.Conv2d、nn.Conv3d、nn.ConvTranspose1d、nn.ConvTranspose2d、nn.ConvTranspose3d が該当します。

weight
- 重み
- init.kaiming_uniform_(a=math.sqrt(5)) で初期化 (He の方法)
bias
- バイアス
- $a = \frac{1}{\sqrt{\text{fan\_in}}}$ としたとき、$[-a, a]$ の一様分布で初期化

全結合層

nn.Linear が該当します。

weight
- 重み
- init.kaiming_uniform_(a=math.sqrt(5)) で初期化 (He の方法)
bias
- バイアス
- $a = \begin{cases} \frac{1}{\sqrt{\text{fan\_in}}} &\text{if } fan\_in > 0 \\ 0 &\text{if } fan\_in \le 0 \end{cases}$ としたとき、$[-a, a]$ の一様分布で初期化

Batch Normalization

nn.BatchNorm1d、nn.BatchNorm2d、nn.BatchNorm3d が該当します。

running_mean
- 平均
- 0で初期化
running_var
- 分散
- 1で初期化
weight
- 重み
- 1で初期化
bias
- バイアス
- 0で初期化

前の記事

Pytorch – DenseNet の仕組みと実装について解説 2021.12.02
次の記事

Pytorch – SqueezeNet の仕組みと実装について解説 2021.12.03