3.4. 線形回帰のスクラッチ実装

ここで、線形回帰の完全に動作する実装に取り組む準備が整った。
この節では、(i) モデル、(ii) 損失関数、(iii) ミニバッチ確率勾配降下法オプティマイザ、(iv) これらすべてをつなぎ合わせる学習関数 を含め、方法全体をスクラッチから実装する。
最後に、 3.3 章 の合成データ生成器を実行し、その結果得られるデータセットに対してモデルを適用する。
現代の深層学習フレームワークはこの作業のほとんどを自動化できるが、スクラッチから実装することこそが、自分が何をしているのかを本当に理解していることを確かめる唯一の方法である。
さらに、モデルをカスタマイズしたり、自前の層や損失関数を定義したりするときには、内部で何が起きているかを理解していることが大いに役立ちる。
この節では、テンソルと自動微分だけに頼りる。
後の節では、以下の構造を保ちながら、深層学習フレームワークの便利な機能を活用した、より簡潔な実装を紹介する。
pytorchmxnetjaxtensorflow
%matplotlib inline
from d2l import torch as d2l
import torch

%matplotlib inline
from d2l import mxnet as d2l
from mxnet import autograd, np, npx
npx.set_np()

%matplotlib inline
from d2l import jax as d2l
from flax import linen as nn
import jax
from jax import numpy as jnp
import optax

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

%matplotlib inline
from d2l import tensorflow as d2l
import tensorflow as tf

3.4.1. モデルの定義

ミニバッチSGDでモデルのパラメータを最適化し始める前に、まずパラメータが必要である。
以下では、平均0、標準偏差0.01の正規分布から乱数を引くことで重みを初期化する。
魔法の数0.01は実践上しばしばうまくいくが、引数 sigma を通して別の値を指定することもできる。
さらに、バイアスは0に設定する。
オブジェクト指向設計においては、コードを d2l.Module のサブクラスの __init__ メソッドに追加する（3.2.2 章 で導入）。
pytorchmxnetjaxtensorflow
class LinearRegressionScratch(d2l.Module):  #@save
    """The linear regression model implemented from scratch."""
    def __init__(self, num_inputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        if tab.selected('mxnet'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1))
            self.b = d2l.zeros(1)
            self.w.attach_grad()
            self.b.attach_grad()
        if tab.selected('pytorch'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1), requires_grad=True)
            self.b = d2l.zeros(1, requires_grad=True)
        if tab.selected('tensorflow'):
            w = tf.random.normal((num_inputs, 1), mean=0, stddev=0.01)
            b = tf.zeros(1)
            self.w = tf.Variable(w, trainable=True)
            self.b = tf.Variable(b, trainable=True)

class LinearRegressionScratch(d2l.Module):  #@save
    """The linear regression model implemented from scratch."""
    def __init__(self, num_inputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        if tab.selected('mxnet'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1))
            self.b = d2l.zeros(1)
            self.w.attach_grad()
            self.b.attach_grad()
        if tab.selected('pytorch'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1), requires_grad=True)
            self.b = d2l.zeros(1, requires_grad=True)
        if tab.selected('tensorflow'):
            w = tf.random.normal((num_inputs, 1), mean=0, stddev=0.01)
            b = tf.zeros(1)
            self.w = tf.Variable(w, trainable=True)
            self.b = tf.Variable(b, trainable=True)

class LinearRegressionScratch(d2l.Module):  #@save
    """The linear regression model implemented from scratch."""
    num_inputs: int
    lr: float
    sigma: float = 0.01

    def setup(self):
        self.w = self.param('w', nn.initializers.normal(self.sigma),
                            (self.num_inputs, 1))
        self.b = self.param('b', nn.initializers.zeros, (1))

class LinearRegressionScratch(d2l.Module):  #@save
    """The linear regression model implemented from scratch."""
    def __init__(self, num_inputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        if tab.selected('mxnet'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1))
            self.b = d2l.zeros(1)
            self.w.attach_grad()
            self.b.attach_grad()
        if tab.selected('pytorch'):
            self.w = d2l.normal(0, sigma, (num_inputs, 1), requires_grad=True)
            self.b = d2l.zeros(1, requires_grad=True)
        if tab.selected('tensorflow'):
            w = tf.random.normal((num_inputs, 1), mean=0, stddev=0.01)
            b = tf.zeros(1)
            self.w = tf.Variable(w, trainable=True)
            self.b = tf.Variable(b, trainable=True)
次に、入力とパラメータを出力に結びつけるモデルを定義する必要がある。
(3.1.4) と同じ記法を用いると、線形モデルでは入力特徴量 \(\mathbf{X}\) とモデルの重み \(\mathbf{w}\) の行列–ベクトル積を取り、各例にオフセット \(b\) を加えるだけです。
積 \(\mathbf{Xw}\) はベクトルであり、\(b\) はスカラーである。
ブロードキャスト機構（2.1.4 章 を参照）により、ベクトルとスカラーを加えると、スカラーはベクトルの各成分に加えられる。
以下の forward メソッドは、add_to_class（3.2.1 章 で導入）を通じて LinearRegressionScratch クラスに登録される。
@d2l.add_to_class(LinearRegressionScratch)  #@save
def forward(self, X):
    return d2l.matmul(X, self.w) + self.b

3.4.2. 損失関数の定義

モデルの更新には損失関数の勾配を取る必要があるので、まず損失関数を定義しておくべきである。
ここでは (3.1.5) の二乗損失関数を使う。
実装では、真の値 y を予測値 y_hat の形状に変換する必要がある。
以下のメソッドが返す結果も y_hat と同じ形状になる。
また、ミニバッチ内の全例にわたる平均損失値も返す。
pytorchmxnetjaxtensorflow
@d2l.add_to_class(LinearRegressionScratch)  #@save
def loss(self, y_hat, y):
    l = (y_hat - y) ** 2 / 2
    return d2l.reduce_mean(l)

@d2l.add_to_class(LinearRegressionScratch)  #@save
def loss(self, y_hat, y):
    l = (y_hat - y) ** 2 / 2
    return d2l.reduce_mean(l)

@d2l.add_to_class(LinearRegressionScratch)  #@save
def loss(self, params, X, y, state):
    y_hat = state.apply_fn({'params': params}, *X)  # X unpacked from a tuple
    l = (y_hat - d2l.reshape(y, y_hat.shape)) ** 2 / 2
    return d2l.reduce_mean(l)

@d2l.add_to_class(LinearRegressionScratch)  #@save
def loss(self, y_hat, y):
    l = (y_hat - y) ** 2 / 2
    return d2l.reduce_mean(l)

3.4.3. 最適化アルゴリズムの定義

3.1 章 で議論したように、線形回帰には閉形式解がある。
しかし、ここでの目的は、より一般的なニューラルネットワークをどのように学習させるかを示すことである。そのためには、ミニバッチSGDの使い方を学ぶ必要がある。
そこで、この機会にSGDの最初の動作例を紹介する。
各ステップで、データセットからランダムに取り出したミニバッチを用いて、パラメータに関する損失の勾配を推定する。
次に、損失を減らす方向にパラメータを更新する。
以下のコードは、パラメータ集合と学習率 lr が与えられたときに更新を適用する。
損失はミニバッチ上の平均として計算されるので、バッチサイズに応じて学習率を調整する必要はない。
後の章では、分散大規模学習で現れるような非常に大きなミニバッチに対して、学習率をどのように調整すべきかを調べる。
今のところは、この依存関係は無視して構わない。
d2l.HyperParameters（3.2.1 章 で導入）のサブクラスとして SGD クラスを定義し、組み込みのSGDオプティマイザと同様のAPIを持たせる。
パラメータは step メソッドで更新する。
zero_grad メソッドはすべての勾配を0に設定し、逆伝播の前に実行しなければならない。
pytorchmxnetjaxtensorflow
class SGD(d2l.HyperParameters):  #@save
    """Minibatch stochastic gradient descent."""
    def __init__(self, params, lr):
        self.save_hyperparameters()

    if tab.selected('mxnet'):
        def step(self, _):
            for param in self.params:
                param -= self.lr * param.grad

    if tab.selected('pytorch'):
        def step(self):
            for param in self.params:
                param -= self.lr * param.grad

        def zero_grad(self):
            for param in self.params:
                if param.grad is not None:
                    param.grad.zero_()

class SGD(d2l.HyperParameters):  #@save
    """Minibatch stochastic gradient descent."""
    def __init__(self, params, lr):
        self.save_hyperparameters()

    if tab.selected('mxnet'):
        def step(self, _):
            for param in self.params:
                param -= self.lr * param.grad

    if tab.selected('pytorch'):
        def step(self):
            for param in self.params:
                param -= self.lr * param.grad

        def zero_grad(self):
            for param in self.params:
                if param.grad is not None:
                    param.grad.zero_()

class SGD(d2l.HyperParameters):  #@save
    """Minibatch stochastic gradient descent."""
    # The key transformation of Optax is the GradientTransformation
    # defined by two methods, the init and the update.
    # The init initializes the state and the update transforms the gradients.
    # https://github.com/deepmind/optax/blob/master/optax/_src/transform.py
    def __init__(self, lr):
        self.save_hyperparameters()

    def init(self, params):
        # Delete unused params
        del params
        return optax.EmptyState

    def update(self, updates, state, params=None):
        del params
        # When state.apply_gradients method is called to update flax's
        # train_state object, it internally calls optax.apply_updates method
        # adding the params to the update equation defined below.
        updates = jax.tree_util.tree_map(lambda g: -self.lr * g, updates)
        return updates, state

    def __call__():
        return optax.GradientTransformation(self.init, self.update)

class SGD(d2l.HyperParameters):  #@save
    """Minibatch stochastic gradient descent."""
    def __init__(self, lr):
        self.save_hyperparameters()

    def apply_gradients(self, grads_and_vars):
        for grad, param in grads_and_vars:
            param.assign_sub(self.lr * grad)

次に、SGD クラスのインスタンスを返す configure_optimizers メソッドを定義する。

@d2l.add_to_class(LinearRegressionScratch)  #@save
def configure_optimizers(self):
    if tab.selected('mxnet') or tab.selected('pytorch'):
        return SGD([self.w, self.b], self.lr)
    if tab.selected('tensorflow', 'jax'):
        return SGD(self.lr)

3.4.4. 学習

これで、すべての部品（パラメータ、損失関数、モデル、オプティマイザ）が揃ったので、主要な学習ループを実装する準備が整った。
このコードを完全に理解することは非常に重要である。というのも、この本で扱う他のすべての深層学習モデルでも、同様の学習ループを使うからである。
各 epoch では、訓練データセット全体を一巡し、すべての例を1回ずつ通過する（例の数がバッチサイズで割り切れると仮定する）。
各 iteration では、訓練例のミニバッチを取り出し、モデルの training_step メソッドを通して損失を計算する。
次に、各パラメータに関する勾配を計算する。
最後に、最適化アルゴリズムを呼び出してモデルパラメータを更新する。
要するに、以下のループを実行する。

  • パラメータ \((\mathbf{w}, b)\) を初期化する

  • 完了するまで繰り返す

    • 勾配 \(\mathbf{g} \leftarrow \partial_{(\mathbf{w},b)} \frac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} l(\mathbf{x}^{(i)}, y^{(i)}, \mathbf{w}, b)\) を計算する

    • パラメータ \((\mathbf{w}, b) \leftarrow (\mathbf{w}, b) - \eta \mathbf{g}\) を更新する

3.3 章 で生成した合成回帰データセットには検証データセットが含まれていないことを思い出そう。
しかし一般には、モデル品質を測定するために検証データセットが欲しくなる。
ここでは、各 epoch ごとに検証データローダーを1回通してモデル性能を測定する。
オブジェクト指向設計に従い、prepare_batch と fit_epoch メソッドは d2l.Trainer クラスに登録される（3.2.4 章 で導入）。
@d2l.add_to_class(d2l.Trainer)  #@save
def prepare_batch(self, batch):
    return batch
pytorchmxnetjaxtensorflow
@d2l.add_to_class(d2l.Trainer)  #@save
def fit_epoch(self):
    self.model.train()
    for batch in self.train_dataloader:
        loss = self.model.training_step(self.prepare_batch(batch))
        self.optim.zero_grad()
        with torch.no_grad():
            loss.backward()
            if self.gradient_clip_val > 0:  # To be discussed later
                self.clip_gradients(self.gradient_clip_val, self.model)
            self.optim.step()
        self.train_batch_idx += 1
    if self.val_dataloader is None:
        return
    self.model.eval()
    for batch in self.val_dataloader:
        with torch.no_grad():
            self.model.validation_step(self.prepare_batch(batch))
        self.val_batch_idx += 1

@d2l.add_to_class(d2l.Trainer)  #@save
def fit_epoch(self):
    for batch in self.train_dataloader:
        with autograd.record():
            loss = self.model.training_step(self.prepare_batch(batch))
        loss.backward()
        if self.gradient_clip_val > 0:
            self.clip_gradients(self.gradient_clip_val, self.model)
        self.optim.step(1)
        self.train_batch_idx += 1
    if self.val_dataloader is None:
        return
    for batch in self.val_dataloader:
        self.model.validation_step(self.prepare_batch(batch))
        self.val_batch_idx += 1

@d2l.add_to_class(d2l.Trainer)  #@save
def fit_epoch(self):
    self.model.training = True
    if self.state.batch_stats:
        # Mutable states will be used later (e.g., for batch norm)
        for batch in self.train_dataloader:
            (_, mutated_vars), grads = self.model.training_step(self.state.params,
                                                           self.prepare_batch(batch),
                                                           self.state)
            self.state = self.state.apply_gradients(grads=grads)
            # Can be ignored for models without Dropout Layers
            self.state = self.state.replace(
                dropout_rng=jax.random.split(self.state.dropout_rng)[0])
            self.state = self.state.replace(batch_stats=mutated_vars['batch_stats'])
            self.train_batch_idx += 1
    else:
        for batch in self.train_dataloader:
            _, grads = self.model.training_step(self.state.params,
                                                self.prepare_batch(batch),
                                                self.state)
            self.state = self.state.apply_gradients(grads=grads)
            # Can be ignored for models without Dropout Layers
            self.state = self.state.replace(
                dropout_rng=jax.random.split(self.state.dropout_rng)[0])
            self.train_batch_idx += 1

    if self.val_dataloader is None:
        return
    self.model.training = False
    for batch in self.val_dataloader:
        self.model.validation_step(self.state.params,
                                   self.prepare_batch(batch),
                                   self.state)
        self.val_batch_idx += 1

@d2l.add_to_class(d2l.Trainer)  #@save
def fit_epoch(self):
    self.model.training = True
    for batch in self.train_dataloader:
        with tf.GradientTape() as tape:
            loss = self.model.training_step(self.prepare_batch(batch))
        grads = tape.gradient(loss, self.model.trainable_variables)
        if self.gradient_clip_val > 0:
            grads = self.clip_gradients(self.gradient_clip_val, grads)
        self.optim.apply_gradients(zip(grads, self.model.trainable_variables))
        self.train_batch_idx += 1
    if self.val_dataloader is None:
        return
    self.model.training = False
    for batch in self.val_dataloader:
        self.model.validation_step(self.prepare_batch(batch))
        self.val_batch_idx += 1
モデルを学習する準備はほぼ整ったが、その前に学習データが必要である。
ここでは SyntheticRegressionData クラスを使い、真のパラメータを与える。
そして、学習率 lr=0.03 でモデルを学習し、max_epochs=3 に設定する。
一般に、epoch数と学習率の両方はハイパーパラメータであることに注意されたい。
ハイパーパラメータの設定は一般に難しく、通常はデータセットを3つに分割したくなる。すなわち、1つは学習用、2つ目はハイパーパラメータ選択用の検証用、そして3つ目は最終評価用である。
ここではこれらの詳細は省きるが、後ほど改めて詳しく取り上げる。
model = LinearRegressionScratch(2, lr=0.03)
data = d2l.SyntheticRegressionData(w=d2l.tensor([2, -3.4]), b=4.2)
trainer = d2l.Trainer(max_epochs=3)
trainer.fit(model, data)
../_images/output_linear-regression-scratch_8b7005_81_0.svg
データセットを自分たちで合成したので、真のパラメータが何であるかを正確に知っている。
したがって、学習ループを通じて真のパラメータと学習したパラメータを比較することで、学習の成功を評価できる。
実際、それらは非常に近い値になる。
pytorchmxnetjaxtensorflow
with torch.no_grad():
    print(f'error in estimating w: {data.w - d2l.reshape(model.w, data.w.shape)}')
    print(f'error in estimating b: {data.b - model.b}')

error in estimating w: tensor([ 0.0709, -0.2553])
error in estimating b: tensor([0.2337])

print(f'error in estimating w: {data.w - d2l.reshape(model.w, data.w.shape)}')
print(f'error in estimating b: {data.b - model.b}')

error in estimating w: [ 0.1114291  -0.12024117]
error in estimating b: [0.1888752]

params = trainer.state.params
print(f"error in estimating w: {data.w - d2l.reshape(params['w'], data.w.shape)}")
print(f"error in estimating b: {data.b - params['b']}")

error in estimating w: [ 0.07124913 -0.19058824]
error in estimating b: [0.23817086]

print(f'error in estimating w: {data.w - d2l.reshape(model.w, data.w.shape)}')
print(f'error in estimating b: {data.b - model.b}')

error in estimating w: [ 0.12973297 -0.16401815]
error in estimating b: [0.25387883]
真のパラメータを正確に復元できることを当然だと思うべきではない。
一般に、深いモデルではパラメータの一意な解は存在せず、線形モデルであっても、ある特徴量が他の特徴量に対して線形従属でない場合にのみ、パラメータを正確に復元できる。
しかし機械学習では、真の潜在パラメータを復元することよりも、高精度な予測をもたらすパラメータのほうが重要であることが多いです (Vapnik, 1992)。
幸いなことに、難しい最適化問題であっても、確率勾配降下法はしばしば驚くほど良い解を見つけられる。これは、深層ネットワークでは高精度な予測をもたらすパラメータの構成が多数存在することが、一因である。

3.4.5. まとめ

この節では、完全に機能するニューラルネットワークモデルと学習ループを実装することで、深層学習システムの設計に向けて大きな一歩を踏み出した。
この過程で、データローダー、モデル、損失関数、最適化手続き、可視化・監視ツールを構築した。
これは、モデル学習に関係するすべての要素を含むPythonオブジェクトを組み立てることで実現した。
これはまだプロフェッショナル品質の実装ではないが、完全に機能し、このようなコードは小さな問題を素早く解くのにすでに役立ちる。
今後の節では、これを より簡潔に（定型コードを避けて）かつ より効率的に（GPUを最大限に活用して）行う方法を見ていく。

3.4.6. 演習

  1. 重みを0で初期化したらどうなるだろうか。アルゴリズムはそれでも動作するだろうか。では、パラメータを0.01ではなく分散1000で初期化したらどうだろうか。

  2. あなたが Georg Simon Ohm で、電圧と電流を関連づける抵抗のモデルを考えようとしているとする。自動微分を使ってモデルのパラメータを学習できるか。

  3. プランクの法則 を使って、スペクトルエネルギー密度から物体の温度を求められるだろうか。参考までに、黒体から放射される放射のスペクトル密度 \(B\) は \(B(\lambda, T) = \frac{2 hc^2}{\lambda^5} \cdot \left(\exp \frac{h c}{\lambda k T} - 1\right)^{-1}\) である。ここで、\(\lambda\) は波長、\(T\) は温度、\(c\) は光速、\(h\) はプランク定数、\(k\) はボルツマン定数である。異なる波長 \(\lambda\) に対するエネルギーを測定したとして、スペクトル密度曲線をプランクの法則にフィットさせる必要がある。

  4. 損失の2階微分を計算したい場合、どのような問題に遭遇する可能性があるか。それらをどう解決するか。

  5. loss 関数で reshape メソッドが必要なのはなぜですか。

  6. 異なる学習率を使って実験し、損失関数の値がどれくらい速く下がるかを調べよ。学習エポック数を増やすことで誤差を減らせますか。

  7. 例の数がバッチサイズで割り切れない場合、epochの終わりに data_iter はどうなるか。

  8. 絶対値損失 (y_hat - d2l.reshape(y, y_hat.shape)).abs().sum() のような別の損失関数を実装してみよ。

    1. 通常のデータで何が起こるか確認せよ。

    2. \(\\mathbf{y}\) のいくつかの要素、たとえば \(y_5 = 10000\) を意図的に摂動させた場合、挙動に違いがあるか確認せよ。

    3. 二乗損失と絶対値損失の長所を組み合わせる安価な解決策を考えられますか。 ヒント: 非常に大きな勾配値をどう避けますか。

  9. なぜデータセットを再シャッフルする必要があるのだろうか。悪意を持って構成されたデータセットが、そうしないと最適化アルゴリズムを破綻させるような例を設計できるか。