6.4. 遅延初期化¶

ここまでのところ、ネットワークの構築においてかなり大雑把にやってもうまくいっているように見えたかもしれない。具体的には、次のような直感に反することを行ってきた。これらは本来うまく動くようには思えないかもしれない。

入力次元を指定せずにネットワークアーキテクチャを定義した。
直前の層の出力次元を指定せずに層を追加した。
さらに、モデルが何個のパラメータを持つべきかを決めるのに十分な情報を与える前に、これらのパラメータを「初期化」した。

コードが実際に動いていることに驚くかもしれない。そもそも、深層学習フレームワークがネットワークの入力次元を知る方法はない。ここでの工夫は、フレームワークが初期化を遅延し、最初にデータをモデルに通すまで待って、その場で各層のサイズを推論することである。

後で畳み込みニューラルネットワークを扱うときには、この手法はさらに便利になる。なぜなら、入力次元（たとえば画像の解像度）が、その後に続く各層の次元に影響するからである。したがって、コードを書く時点では次元の値を知らなくてもパラメータを設定できる能力は、モデルの指定やその後の修正を大幅に簡単にしてくれる。それでは、初期化の仕組みをさらに詳しく見ていこう。

pytorch mxnet jax tensorflow

from d2l import torch as d2l
import torch
from torch import nn

from mxnet import np, npx
from mxnet.gluon import nn
npx.set_np()

from d2l import jax as d2l
from flax import linen as nn
import jax
from jax import numpy as jnp

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

import tensorflow as tf

まず、MLP をインスタンス化してみよう。

pytorch mxnet jax tensorflow

net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))

net = nn.Sequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))

net = nn.Sequential([nn.Dense(256), nn.relu, nn.Dense(10)])

net = tf.keras.models.Sequential([
    tf.keras.layers.Dense(256, activation=tf.nn.relu),
    tf.keras.layers.Dense(10),
])

この時点では、入力次元がまだ不明なので、ネットワークは入力層の重みの次元を知ることができない。

したがって、フレームワークはまだどのパラメータも初期化していない。以下でパラメータにアクセスしようとして確認してみよう。

pytorch mxnet jax tensorflow

net[0].weight

<UninitializedParameter>

print(net.collect_params)
print(net.collect_params())

<bound method Block.collect_params of Sequential(
  (0): Dense(-1 -> 256, Activation(relu))
  (1): Dense(-1 -> 10, linear)
)>
sequential0_ (
  Parameter dense0_weight (shape=(256, -1), dtype=float32)
  Parameter dense0_bias (shape=(256,), dtype=float32)
  Parameter dense1_weight (shape=(10, -1), dtype=float32)
  Parameter dense1_bias (shape=(10,), dtype=float32)
)

params = net.init(d2l.get_key(), jnp.zeros((2, 20)))
jax.tree_util.tree_map(lambda x: x.shape, params).tree_flatten_with_keys()

(((DictKey(key='params'),
   {'layers_0': {'bias': (256,), 'kernel': (20, 256)},
    'layers_2': {'bias': (10,), 'kernel': (256, 10)}}),),
 ('params',))

[net.layers[i].get_weights() for i in range(len(net.layers))]

[[], []]

次に、ネットワークにデータを通して、フレームワークにようやくパラメータを初期化させてみよう。

pytorch mxnet jax tensorflow

X = torch.rand(2, 20)
net(X)

net[0].weight.shape

torch.Size([256, 20])

net.initialize()
net.collect_params()

[07:07:50] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU

sequential0_ (
  Parameter dense0_weight (shape=(256, -1), dtype=float32)
  Parameter dense0_bias (shape=(256,), dtype=float32)
  Parameter dense1_weight (shape=(10, -1), dtype=float32)
  Parameter dense1_bias (shape=(10,), dtype=float32)
)

@d2l.add_to_class(d2l.Module)  #@save
def apply_init(self, dummy_input, key):
    params = self.init(key, *dummy_input)  # dummy_input tuple unpacked
    return params

X = tf.random.uniform((2, 20))
net(X)
[w.shape for w in net.get_weights()]

[(20, 256), (256,), (256, 10), (10,)]

入力次元 20 が分かればすぐに、フレームワークは 20 の値を代入することで最初の層の重み行列の形状を特定できる。最初の層の形状が分かると、フレームワークは次の層へ進み、計算グラフに沿って順に処理し、すべての形状が分かるまで続ける。この場合、遅延初期化が必要なのは最初の層だけだが、フレームワークは順次に初期化を行う。すべてのパラメータ形状が分かると、フレームワークはようやくパラメータを初期化できる。

次のメソッドは、ダミー入力をネットワークに通して予備実行を行い、すべてのパラメータ形状を推論したうえでパラメータを初期化する。これは、デフォルトのランダム初期化を望まない場合に後で使われる。

pytorch mxnet

@d2l.add_to_class(d2l.Module)  #@save
def apply_init(self, inputs, init=None):
    self.forward(*inputs)
    if init is not None:
        self.net.apply(init)

X = np.random.uniform(size=(2, 20))
net(X)

net.collect_params()

sequential0_ (
  Parameter dense0_weight (shape=(256, 20), dtype=float32)
  Parameter dense0_bias (shape=(256,), dtype=float32)
  Parameter dense1_weight (shape=(10, 256), dtype=float32)
  Parameter dense1_bias (shape=(10,), dtype=float32)
)

6.4.1. 要約¶

遅延初期化は便利である。フレームワークがパラメータ形状を自動的に推論できるため、アーキテクチャの修正が容易になり、よくあるエラーの原因を一つ取り除ける。モデルにデータを通すことで、フレームワークにようやくパラメータを初期化させることができる。

6.4.2. 演習¶

最初の層には入力次元を指定するが、その後の層には指定しない場合、どうなるか？すぐに初期化されるか？
次元が一致しないように指定した場合、どうなるか？
入力の次元が変化する場合、何をする必要があるか？ヒント: パラメータ共有を見てみよう。