4.4. スクラッチからのソフトマックス回帰の実装¶

ソフトマックス回帰は非常に基本的な手法なので、ぜひ自分で実装できるようになっておくべきだと考える。ここでは、モデルのソフトマックス固有の部分の定義に限定し、線形回帰の節で使った他の構成要素、たとえば学習ループなどは再利用する。

from d2l import torch as d2l
import torch

from d2l import mxnet as d2l
from mxnet import autograd, np, npx, gluon
npx.set_np()

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

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

from d2l import tensorflow as d2l
import tensorflow as tf

4.4.1. ソフトマックス¶

まず最も重要な部分から始めよう。それは、スカラーを確率へ写像することである。復習として、 2.3.6 章および 2.3.7 章で説明したように、テンソルの特定の次元に沿った和演算を思い出そう。行列 X に対して、すべての要素の和を計算する（デフォルト）ことも、特定の軸に沿った要素の和だけを計算することもできる。 axis 変数を使うと、行方向および列方向の和を計算できる。

X = d2l.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
d2l.reduce_sum(X, 0, keepdims=True), d2l.reduce_sum(X, 1, keepdims=True)

(tensor([[5., 7., 9.]]),
 tensor([[ 6.],
         [15.]]))

ソフトマックスの計算には3つの手順がある。 (i) 各要素の指数を取る。 (ii) 各行について和を取り、各サンプルの正規化定数を計算する。 (iii) 各行をその正規化定数で割り、結果の和が1になるようにする。

(4.4.1)¶\[\mathrm{softmax}(\mathbf{X})_{ij} = \frac{\exp(\mathbf{X}_{ij})}{\sum_k \exp(\mathbf{X}_{ik})}.\]

分母の（対数を取ったもの）は（対数）分配関数と呼ばれる。これは熱力学的アンサンブルにおけるすべての可能な状態を総和するために、統計物理学で導入された。実装は簡単である。

def softmax(X):
    X_exp = d2l.exp(X)
    partition = d2l.reduce_sum(X_exp, 1, keepdims=True)
    return X_exp / partition  # ここではブロードキャスト機構が適用される

任意の入力 X に対して、各要素を非負の数に変換する。各行の和は1になる。これは確率として必要な性質である。注意: 上のコードは、非常に大きい値や非常に小さい値に対しては頑健ではない。何が起きているかを示すには十分であるが、実用上の目的でこのコードをそのまま使うべきではない。深層学習フレームワークにはこのような保護機構が組み込まれており、今後は組み込みのソフトマックスを使う。

pytorch mxnet jax tensorflow

X = d2l.rand((2, 5))
X_prob = softmax(X)
X_prob, d2l.reduce_sum(X_prob, 1)

(tensor([[0.1481, 0.2050, 0.2577, 0.1232, 0.2661],
         [0.2167, 0.2010, 0.2468, 0.1128, 0.2227]]),
 tensor([1.0000, 1.0000]))

X = d2l.rand(2, 5)
X_prob = softmax(X)
X_prob, d2l.reduce_sum(X_prob, 1)

(array([[0.17777154, 0.1857739 , 0.20995119, 0.23887765, 0.18762572],
        [0.24042214, 0.1757977 , 0.23786479, 0.15572716, 0.19018826]]),
 array([1., 1.]))

X = jax.random.uniform(jax.random.PRNGKey(d2l.get_seed()), (2, 5))
X_prob = softmax(X)
X_prob, d2l.reduce_sum(X_prob, 1)

(Array([[0.16229036, 0.13327661, 0.22257593, 0.1909562 , 0.2909009 ],
        [0.30952108, 0.17390645, 0.18135622, 0.19830671, 0.13690951]],      dtype=float32),
 Array([1., 1.], dtype=float32))

X = d2l.rand((2, 5))
X_prob = softmax(X)
X_prob, d2l.reduce_sum(X_prob, 1)

(<tf.Tensor: shape=(2, 5), dtype=float32, numpy=
 array([[0.30452266, 0.15993236, 0.17174596, 0.160699  , 0.20310006],
        [0.24015403, 0.09632394, 0.21000664, 0.25052655, 0.20298874]],
       dtype=float32)>,
 <tf.Tensor: shape=(2,), dtype=float32, numpy=array([1.       , 0.9999999], dtype=float32)>)

4.4.2. モデル¶

これで、ソフトマックス回帰モデルを実装するために必要なものはすべて揃った。線形回帰の例と同様に、各データ例は固定長ベクトルで表現される。ここでの生データは \(28 \times 28\) ピクセルの画像からなるので、各画像を平坦化し、長さ784のベクトルとして扱いる。後の章では、空間構造をより自然に活用できる畳み込みニューラルネットワークを紹介する。

ソフトマックス回帰では、ネットワークの出力数はクラス数と等しくなければならない。データセットには10クラスあるので、ネットワークの出力次元は10である。したがって、重みは \(784 \times 10\) の行列と、バイアス用の \(1 \times 10\) の行ベクトルから構成される。線形回帰と同様に、重み W はガウスノイズで初期化する。バイアスはゼロで初期化する。

pytorch mxnet jax tensorflow

class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = torch.normal(0, sigma, size=(num_inputs, num_outputs),
                              requires_grad=True)
        self.b = torch.zeros(num_outputs, requires_grad=True)

    def parameters(self):
        return [self.W, self.b]

class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = np.random.normal(0, sigma, (num_inputs, num_outputs))
        self.b = np.zeros(num_outputs)
        self.W.attach_grad()
        self.b.attach_grad()

    def collect_params(self):
        return [self.W, self.b]

class SoftmaxRegressionScratch(d2l.Classifier):
    num_inputs: int
    num_outputs: int
    lr: float
    sigma: float = 0.01

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

class SoftmaxRegressionScratch(d2l.Classifier):
    def __init__(self, num_inputs, num_outputs, lr, sigma=0.01):
        super().__init__()
        self.save_hyperparameters()
        self.W = tf.random.normal((num_inputs, num_outputs), 0, sigma)
        self.b = tf.zeros(num_outputs)
        self.W = tf.Variable(self.W)
        self.b = tf.Variable(self.b)

以下のコードでは、ネットワークが各入力をどのように出力へ写像するかを定義する。バッチ内の各 \(28 \times 28\) ピクセル画像は、モデルに通す前に reshape を使ってベクトルに平坦化していることに注意されたい。

@d2l.add_to_class(SoftmaxRegressionScratch)
def forward(self, X):
    X = d2l.reshape(X, (-1, self.W.shape[0]))
    return softmax(d2l.matmul(X, self.W) + self.b)

4.4.3. 交差エントロピー損失¶

次に、交差エントロピー損失関数を実装する必要がある (4.1.2 章で導入した)。これは深層学習全体の中でも最も一般的な損失関数かもしれない。現時点では、深層学習の応用のうち、回帰問題よりも分類問題として自然に定式化できるものの方がはるかに多くある。

交差エントロピーは、真のラベルに割り当てられた予測確率の負の対数尤度を取ることを思い出そう。効率のため、Python の for ループは避け、代わりにインデックス参照を使う。特に、\(\mathbf{y}\) の one-hot エンコーディングにより、 \(\hat{\mathbf{y}}\) の対応する項を選択できる。

これを確認するために、3クラスに対する予測確率を持つ2つのデータ例 y_hat と、それに対応するラベル y を作成する。正しいラベルはそれぞれ \(0\) と \(2\)（すなわち第1クラスと第3クラス）である。 y を y_hat 内の確率のインデックスとして使うことで、項を効率よく取り出せる。

pytorch mxnet jax tensorflow

y = d2l.tensor([0, 2])
y_hat = d2l.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]

tensor([0.1000, 0.5000])

y = d2l.tensor([0, 2])
y_hat = d2l.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]

array([0.1, 0.5])

y = d2l.tensor([0, 2])
y_hat = d2l.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]

Array([0.1, 0.5], dtype=float32)

y_hat = tf.constant([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y = tf.constant([0, 2])
tf.boolean_mask(y_hat, tf.one_hot(y, depth=y_hat.shape[-1]))

<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.1, 0.5], dtype=float32)>

これで、選択した確率の対数を平均することで、交差エントロピー損失関数を実装できる。

pytorch mxnet jax tensorflow

def cross_entropy(y_hat, y):
    return -d2l.reduce_mean(d2l.log(y_hat[list(range(len(y_hat))), y]))

cross_entropy(y_hat, y)

tensor(1.4979)

def cross_entropy(y_hat, y):
    return -d2l.reduce_mean(d2l.log(y_hat[list(range(len(y_hat))), y]))

cross_entropy(y_hat, y)

array(1.4978662)

def cross_entropy(y_hat, y):
    return -d2l.reduce_mean(d2l.log(y_hat[list(range(len(y_hat))), y]))

cross_entropy(y_hat, y)

Array(1.4978662, dtype=float32)

def cross_entropy(y_hat, y):
    return -tf.reduce_mean(tf.math.log(tf.boolean_mask(
        y_hat, tf.one_hot(y, depth=y_hat.shape[-1]))))

cross_entropy(y_hat, y)

<tf.Tensor: shape=(), dtype=float32, numpy=1.4978662>

pytorch mxnet jax tensorflow

@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)

@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)

@d2l.add_to_class(SoftmaxRegressionScratch)
@partial(jax.jit, static_argnums=(0))
def loss(self, params, X, y, state):
    def cross_entropy(y_hat, y):
        return -d2l.reduce_mean(d2l.log(y_hat[list(range(len(y_hat))), y]))
    y_hat = state.apply_fn({'params': params}, *X)
    # 返される空の辞書は補助データのプレースホルダであり、
    # 後で（たとえば batch norm のために）使われます
    return cross_entropy(y_hat, y), {}

@d2l.add_to_class(SoftmaxRegressionScratch)
def loss(self, y_hat, y):
    return cross_entropy(y_hat, y)

4.4.4. 学習¶

3.4 章で定義した fit メソッドを再利用して、10エポックでモデルを学習する。エポック数（max_epochs）、ミニバッチサイズ（batch_size）、学習率（lr）は調整可能なハイパーパラメータであることに注意されたい。つまり、これらの値は主たる学習ループの中で学習されるわけではないが、学習性能と汎化性能の両方において、モデルの性能に影響を与える。実際には、データの検証分割に基づいてこれらの値を選び、最終的には テスト 分割で最終モデルを評価したいだろう。 3.6.3 章で述べたように、 Fashion-MNIST のテストデータを検証セットとして扱い、この分割上で検証損失と検証精度を報告する。

data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegressionScratch(num_inputs=784, num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)

../_images/output_softmax-regression-scratch_009b7c_96_0.svg

4.4.5. 予測¶

学習が完了したので、モデルはいくつかの画像を分類する準備ができた。

X, y = next(iter(data.val_dataloader()))
if tab.selected('pytorch', 'mxnet', 'tensorflow'):
    preds = d2l.argmax(model(X), axis=1)
if tab.selected('jax'):
    preds = d2l.argmax(model.apply({'params': trainer.state.params}, X), axis=1)
preds.shape

torch.Size([256])

私たちがより関心を持つのは、誤ってラベル付けされた画像である。実際のラベル（テキスト出力の1行目）とモデルの予測（テキスト出力の2行目）を比較することで、それらを可視化する。

wrong = d2l.astype(preds, y.dtype) != y
X, y, preds = X[wrong], y[wrong], preds[wrong]
labels = [a+'\n'+b for a, b in zip(
    data.text_labels(y), data.text_labels(preds))]
data.visualize([X, y], labels=labels)

../_images/output_softmax-regression-scratch_009b7c_100_0.svg

4.4.6. まとめ¶

ここまでで、線形回帰と分類問題を解く経験を少し積んできた。これにより、統計モデリングの1960〜1970年代におけるいわば最先端に到達したと言ってよいだろう。次の節では、このモデルを深層学習フレームワークを活用してはるかに効率よく実装する方法を示す。

4.4.7. 演習¶

この節では、ソフトマックス演算の数学的定義に基づいてソフトマックス関数を直接実装した。 4.1 章で述べたように、これは数値的不安定性を引き起こす可能性がある。
1. 入力に値 \(100\) が含まれていても softmax が正しく動作するか確認せよ。
2. すべての入力の最大値が \(-100\) より小さい場合でも softmax が正しく動作するか確認せよ。
3. 引数の最大要素に対する相対値を見ることで修正を実装せよ。
交差エントロピー損失関数 \(\sum_i y_i \log \hat{y}_i\) の定義に従う cross_entropy 関数を実装せよ。
1. この節のコード例で試してみよ。
2. なぜより遅く動作すると考えられるか。
3. それを使うべきか。どのような場合に使うのが理にかなっているか。
4. 何に注意する必要があるか。ヒント: 対数の定義域を考えよ。
最もありそうなラベルを返すことは、常に良い考えだろうか。たとえば、医療診断でこれを行うだろうか。どのように対処しようとするか。
いくつかの特徴量に基づいて次の単語を予測するためにソフトマックス回帰を使いたいと仮定する。大きな語彙から生じる問題にはどのようなものがあるか。
この節のコードのハイパーパラメータをいろいろ試してみよ。特に:
1. 学習率を変えたときに検証損失がどのように変化するかをプロットせよ。
2. ミニバッチサイズを変えると検証損失と学習損失は変化するか。効果が見えるまでに、どれくらい大きくまたは小さくする必要があるか。