PyTorchで作る最初のニューラルネットワーク

レッスン04では、numpyで学習ループを手作りしました。順伝播、損失、勾配、勾配と逆向きのステップです。本レッスンでは同じことをPyTorchで行います。勾配は自動で計算され、それを使って手書き数字を読む小さなニューラルネットワークを学習します。スクリプト全体はノートPCのCPUで5秒以内に実行できます。scikit-learnのfitとpredictを使ったことがあり、モデルがネットワークのときに箱の中で何が起きているのかを見たい人のためのレッスンです。

データセットはsklearn.datasets.load_digitsです。8 x 8ピクセルの画像1,797枚で、それぞれが0から9の数字、各ピクセルは0から16のグレーレベルです。scikit-learnに同梱されているので何もダウンロードしません。画像を数字の格子として表示して数字を目で認識できるほど小さく、大きなデータセットでは不可能な有用な健全性チェックができます。

テンソルと自動微分

テンソルはPyTorchの配列です。numpy配列のように振る舞いますが、2つの追加があります。アクセラレータ（GPU、またはAppleのMPSバックエンド）に置けること、そして適用された演算を記録して後から勾配を計算できることです。この2つ目の機能であるautogradのおかげで、レッスン04の手書きの勾配コードは不要になります。

import torch

torch.manual_seed(0)
a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])   # shape (2, 3)
w = torch.randn(3, 4)                                   # shape (3, 4)
b = torch.zeros(4)
out = a @ w + b                                         # matrix product, shape (2, 4)
print("a.shape =", tuple(a.shape), " w.shape =", tuple(w.shape), " out.shape =", tuple(out.shape))
print("dtype:", out.dtype, " device:", out.device)

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print("dy/dx at x=2:", x.grad.item())   # 2x + 3 = 7

a.shape = (2, 3)  w.shape = (3, 4)  out.shape = (2, 4)
dtype: torch.float32  device: cpu
dy/dx at x=2: 7.0

前半は、これから絶えず行うことになる形状の計算です。3特徴量の2行のバッチに3 x 4の重み行列を掛けると、4出力の2行になります。バイアスbは形状(4,)で、ブロードキャストによって各行に加えられます。規則の全体はテンソル形状のリファレンスにあります。

後半がautogradです。requires_grad=Trueを設定すると、PyTorchはxが関わるすべての演算を記録します。スカラーyに対してbackward()を呼ぶと、その記録を逆向きにたどり、x.gradにdy/dxを書き込みます。ネットワークではxが数千の重みに、yが損失に置き換わり、同じ呼び出し1回ですべての勾配が一度に埋まります。

数字データをテンソルとして読み込む

データの準備は通常のscikit-learnです。層化分割を行い、学習セットだけの統計量で標準化します（レッスン02の規則）。ここでの分割をテストではなく検証と呼ぶのは、学習中にこれを見ながらいつ止めるか、どの設定を残すかを決めるからで、つまり判断に使われています。本当のテストセットはさらに別に取り置くことになります。

import numpy as np
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X, y = digits.data.astype(np.float32), digits.target
print("images:", X.shape, " pixel range:", X.min(), "to", X.max())

X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=0
)
mean, std = X_train.mean(), X_train.std()
X_train = (X_train - mean) / std      # scale with training statistics only
X_val = (X_val - mean) / std

X_train_t = torch.from_numpy(X_train)
y_train_t = torch.from_numpy(y_train)
X_val_t = torch.from_numpy(X_val)
y_val_t = torch.from_numpy(y_val)
print("train:", tuple(X_train_t.shape), "val:", tuple(X_val_t.shape), "labels dtype:", y_train_t.dtype)

images: (1797, 64)  pixel range: 0.0 to 16.0
train: (1437, 64) val: (360, 64) labels dtype: torch.int64

後でエラーメッセージに悩まされずに済む、dtypeに関する2つの細部があります。特徴量はfloat32です。これはPyTorchのデフォルトであり、nn.Linearが期待するものです（numpyのデフォルトはfloat64なのでastypeが必要です）。分類のラベルは整数のint64です。損失関数がこれを出力列へのインデックスとして使うからです。各画像の64ピクセルは1行に平坦化されるので、ネットワークが見るのは格子ではなくベクトルです。格子を保つ畳み込みネットワークはMNIST CNNプロジェクトの主題です。

nn.Moduleとしてのモデル

多層パーセプトロン（MLP）は、線形層の間に非線形性を挟んで積み重ねたものです。隠れ層が1つなら、h = relu(x @ W1 + b1)を計算し、次にscores = h @ W2 + b2を計算します。図は本レッスンのネットワークの形状を示しています。

ReLUがなければ、2つの線形層は1つに潰れてしまい（行列の積は行列です）、モデルはレッスン03のロジスティック回帰を遠回りして作っただけになります。ReLU、つまりmax(0, h)は、この潰れを防ぐ最も単純な関数です。負の活性化をゼロにし、正の活性化をそのまま通し、勾配は0か1なので学習が安定します。

出力はロジットと呼ばれます。クラスごとに1つの、生の、上下限のないスコアです。確率ではありません。交差エントロピー損失は内部でソフトマックスを適用して確率に変換し、正解クラスに割り当てられた確率の負の対数を取ります。ソフトマックスをモデル内ではなく損失の内部に置くほうが数値的に安全であり、あらゆるPyTorchチュートリアルが従う慣習でもあるので、分類器の末尾にソフトマックス層を追加しないでください。

PyTorchでは、モデルはnn.Moduleを継承するクラスで、__init__で層を構築し、forwardで順伝播を定義します。モジュールはselfに代入した層の中のすべてのパラメータを追跡するので、オプティマイザはmodel.parameters()でそれらを見つけられます。

学習ループ

下のブロックが全体です。モデル、損失、オプティマイザ、ループ。ループ本体をレッスン04の5つのステップと照らし合わせて読んでください。順伝播（logits = model(...)）、損失、古い勾配のゼロ化、逆伝播、ステップです。唯一の新しい要素はミニバッチ処理です。各エポックで学習行をシャッフルし、32行ずつ進むので、1エポックあたりの更新は1回ではなく45回になります。

import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
np.random.seed(0)

digits = load_digits()
X, y = digits.data.astype(np.float32), digits.target
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
mean, std = X_train.mean(), X_train.std()
X_train_t = torch.from_numpy((X_train - mean) / std)
X_val_t = torch.from_numpy((X_val - mean) / std)
y_train_t = torch.from_numpy(y_train)
y_val_t = torch.from_numpy(y_val)


class Mlp(nn.Module):
    def __init__(self, n_in: int = 64, n_hidden: int = 64, n_out: int = 10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, n_hidden),
            nn.ReLU(),
            nn.Linear(n_hidden, n_out),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)          # raw scores (logits), one per class


model = Mlp()
n_params = sum(p.numel() for p in model.parameters())
print("parameters:", n_params)

loss_fn = nn.CrossEntropyLoss()
optimiser = torch.optim.SGD(model.parameters(), lr=0.05)
batch_size = 32
n_train = len(X_train_t)

for epoch in range(1, 21):
    model.train()
    perm = torch.randperm(n_train)
    running = 0.0
    for start in range(0, n_train, batch_size):
        idx = perm[start:start + batch_size]
        logits = model(X_train_t[idx])
        loss = loss_fn(logits, y_train_t[idx])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
        running += loss.item() * len(idx)
    train_loss = running / n_train

    model.eval()
    with torch.no_grad():
        val_logits = model(X_val_t)
        val_loss = loss_fn(val_logits, y_val_t).item()
        val_acc = (val_logits.argmax(dim=1) == y_val_t).float().mean().item()
    print(f"epoch {epoch:2d}  train loss {train_loss:.3f}  val loss {val_loss:.3f}  val acc {val_acc:.3f}")

parameters: 4810
epoch  1  train loss 1.895  val loss 1.382  val acc 0.850
epoch  2  train loss 0.966  val loss 0.685  val acc 0.906
epoch  3  train loss 0.519  val loss 0.458  val acc 0.914
epoch  4  train loss 0.357  val loss 0.345  val acc 0.936
epoch  5  train loss 0.276  val loss 0.287  val acc 0.950
epoch  6  train loss 0.231  val loss 0.252  val acc 0.950
epoch  7  train loss 0.198  val loss 0.236  val acc 0.956
epoch  8  train loss 0.175  val loss 0.217  val acc 0.953
epoch  9  train loss 0.161  val loss 0.198  val acc 0.956
epoch 10  train loss 0.145  val loss 0.180  val acc 0.961
epoch 11  train loss 0.135  val loss 0.176  val acc 0.958
epoch 12  train loss 0.123  val loss 0.163  val acc 0.969
epoch 13  train loss 0.117  val loss 0.167  val acc 0.961
epoch 14  train loss 0.109  val loss 0.161  val acc 0.961
epoch 15  train loss 0.104  val loss 0.154  val acc 0.964
epoch 16  train loss 0.098  val loss 0.148  val acc 0.961
epoch 17  train loss 0.093  val loss 0.141  val acc 0.972
epoch 18  train loss 0.089  val loss 0.139  val acc 0.967
epoch 19  train loss 0.085  val loss 0.144  val acc 0.961
epoch 20  train loss 0.081  val loss 0.136  val acc 0.969

パラメータ数は、第1層が64 x 64 + 64、第2層が64 x 10 + 10で、合計4,810です。各エポックは45ミニバッチで、全体の実行はCPUで約2秒です。シードを設定していても、浮動小数点のリダクションはハードウェア間でビット単位では一致しないため、正確な損失値はマシンやPyTorchのバージョンによってわずかに異なりますが、曲線の形は同じになります。

曲線を読む

両方の損失は5エポックほど急激に下がり、その後平坦になります。検証損失はエポック4あたりから学習損失の上に位置し、差はゆっくり広がります。モデルは未知の行より学習行のほうを少しだけよく当てはめており、これは正常で、ここでは小さな差です。学習損失が下がり続けているのに検証損失が上がり始めたなら、それは過学習であり、対処は早めに止める、隠れ層を縮める、重み減衰を加える、データを増やす、のいずれかです。このデータでこのネットワークを20エポック回しても、そこまでは行きません。検証の正解率は97%近くで終わります。これが何を示しているかについては正直でいてください。同じ標準化済み分割に対するscikit-learnのロジスティック回帰も約97%を出します。この数字データは小さく、きれいで、ほぼ線形分離可能だからです。ここでは隠れ層は正解率を買ってくれません。本レッスンが買うのは仕組みのほうで、それはデータが224 x 224ピクセルの画像になり、モデルが100層になって線形モデルがまったく歯が立たなくなっても変わりません。

エポック1の検証損失が学習損失より低いことに驚く人がいます。学習損失は重みがまだ動いている最中にエポック全体で平均されるので、序盤の悪いバッチを含んでいます。検証損失はエポックの終わりに、改善された重みで1回だけ測定されます。

ループの中には忘れやすい行が3つあり、それぞれエラーではなく静かなバグを生みます。optimiser.zero_grad()が必要なのは、PyTorchが勾配を.gradに置き換えるのではなく累積するからで、これがなければ毎回のステップでそれまでの全勾配の合計が使われます。model.eval()とmodel.train()は、ドロップアウトやバッチ正規化などの層の挙動を切り替えます。このネットワークにはどちらもないので今は何もしませんが、初めてそれらを追加するときにこの習慣が効いてきます。torch.no_grad()はautogradに検証の順伝播を記録しないよう伝え、メモリと時間を節約します。PyTorch学習ループのリファレンスに、これらとその他の標準的な行を1か所にまとめています。

規模が大きくなると何が変わるか

本レッスンのすべては、形を変えずにスケールアップします。より大きなネットワークはnn.Sequentialの層が増えます。より大きなデータセットは、ここで使った手動のrandpermスライスの代わりに、ミニバッチを生成するDataLoaderで包みます。GPUはモデルと各バッチに.to(device)を呼ぶことで使います。現代の仕事の大半ではSGDの代わりにtorch.optim.Adamを使います。パラメータごとにステップ幅を適応させるためで、代償は少しのメモリです。ループ本体はそれ以外同じ5行であり、見るべき曲線も同じです。続きは深層学習ハブとPyTorchプラットフォームページで扱います。

演習

1. n_hiddenを8に、次に256に変え、それぞれ20エポック実行してください。3つのサイズすべてについて最終的な学習損失と検証損失を記録してください。両者の差が最も大きいのはどれですか。最小のネットワークは学習不足（underfitting）になっていますか。
2. torch.optim.SGD(..., lr=0.05)をtorch.optim.Adam(model.parameters(), lr=1e-3)に置き換え、5エポック後の曲線を比較してください。次にSGDでlr=0.5を試し、損失に何が起きるか記述してください。
3. 学習後、モデルが間違えた検証画像を見つけてください。wrong = (val_logits.argmax(dim=1) != y_val_t).nonzero()です。そのうち1枚をX_val[i].reshape(8, 8).round(1)で8 x 8の格子として表示し、自分なら正解できたか判断してください。

まとめ

• テンソルは、アクセラレータに置くことができ、演算を記録してbackward()で自動的に勾配を計算できる配列です。
• nn.Moduleは__init__で層を、forwardで計算を宣言します。層の中のパラメータは自動で追跡されます。
• 多層パーセプトロンは、線形、ReLU、線形です。非線形性があるからこそ単一の線形モデル以上のものになります。出力はロジットで、交差エントロピー損失が自身でソフトマックスを適用します。
• ループは、順伝播、損失、zero_grad、backward、stepで、シャッフルしたミニバッチに対して回し、各エポックの終わりにno_gradの検証パスを行います。
• 学習損失と検証損失を一緒に見てください。検証損失が上昇しながら差が広がるのは過学習です。損失がln(クラス数)で止まっているなら、何も学習されていません。

さらに読む

• PyTorch documentation, torch.nn: Module、Linear、ReLU、CrossEntropyLoss、その他の層のカタログ。
• PyTorch documentation, Autograd mechanics: 演算の記録がどう構築され、backward()がそれをどうたどるか。
• PyTorch documentation, torch.optim: SGD、Adam、学習率スケジューラ。
• Rumelhart, Hinton and Williams, Learning representations by back-propagating errors, Nature, 1986: 多層ネットワークの学習を実用的にした論文。