AI/ML Field Guide
English
プロジェクト

PyTorchでCPU上に作るMNIST級のCNN

16x16の数字画像で2層の畳み込みネットワークをノートPCのCPUで1秒以内に学習し、損失曲線を読み、混同行列を出力して、どの数字を間違えるかを調べます。

入門PyTorch更新 2026年11月13日
実行環境ノートPCのCPUまたはApple Silicon、5分以内
所要時間約2時間
使用ライブラリpytorch

作るもの

このプロジェクトは、誰もが一度は書く畳み込みネットワークです。畳み込み層2つ、プーリング層2つ、全結合層2つ、出力クラス10個。データセットはscikit-learnに同梱のload_digitsで、8x8ピクセルの手書き数字1,797枚を、畳み込みに仕事を与えるために16x16にアップサンプリングします。ダウンロードもtorchvisionも不要で、学習全体はノートPCのCPUで約1秒です。MNIST本体は28x28の画像70,000枚でダウンロードが必要なので、後で最初に試すものとして残しています。モデルで変わるのは数値1つだけです。

データローダー、モデル、学習ループ、評価を手書きし、エポックごとに損失とテスト精度を出力し、重みを保存し、別のプログラムで再読み込みして混同行列を出力します。面白いのは混同行列です。精度の数値1つは、モデルが95%正しいと言うだけです。行列は、どの数字をどの数字と取り違え続けているかを示し、それが次に何をすべきかを教えてくれます。

これはPyTorchで作る最初のニューラルネットワークと学習ループのレッスンの続きです。PyTorch学習ループのチートシートにはループの骨格が1ページにまとまっています。

必要なもの

  • Python 3.10以降とtorch（CPUビルドで十分）、numpy、scikit-learn。当サイトではPyTorch 2.14.1とscikit-learn 1.9.1でテストしました。
  • GPUは不要です。2回の実行で同じ数値が出るように、コードはdevice = torch.device("cpu")に固定しています。「試してみること」の節でApple SiliconやCUDA向けの1行の変更を示します。
  • データとモデルに約40 MBのメモリ。保存される重みは155 kBです。

2つのPythonブロックはどちらも、Apple SiliconのMac mini上で、Field Guideのパイプライン環境を使い、空のディレクトリから実行しました。所要時間は各ブロックに記載しています。

設計

入力1 x 16 x 16conv1 3x316 x 16 x 16pool16 x 8 x 8conv2 3x332 x 8 x 8pool32 x 4 x 4flatten512fc1 + ReLU64fc210 logitsCE損失
16x16の画像1枚がネットワークを通る際の形状。各段階でチャネル x 高さ x 幅。

なぜ8x8を16x16にアップサンプリングするのか。 2x2の最大プーリングを2回行うと、空間方向の各次元は4分の1になります。8x8からでは2x2のマップしか残らず、画像とはほとんど言えません。16x16からなら4x4が残り、2つ目の畳み込みが単一ピクセルではなくストロークを見るのに十分です。バイリニア補間は情報を追加せず、カーネルに余地を与えるだけです。同じコードをsize=28とMNIST入力で動かすには、fc1の入力サイズを32 * 7 * 7に変えるだけで済みます。

なぜこの層構成なのか。 conv1は3x3ピクセルのフィルタを16個学習します。エッジやストロークの端です。conv2はそれらをより広い視野で組み合わせたものを32個学習します。パディング1により各畳み込みを通しても空間サイズが保たれ、縮小するのはプーリングだけになります。各ReLUの後にプーリングが続き、2x2の各ウィンドウで最も強い応答を残して、次の層の仕事を半分にします。分類器は64ユニットの隠れ層1つです。全体で38,282パラメータで、学習ブロックがこの数を出力します。

なぜDataLoaderではなく素のループなのか。 データセットは1,437枚の学習画像で、1つのテンソルに収まります。torch.randpermでシャッフルして32個ずつのバッチにスライスするほうが、このサイズではDataLoaderより読みやすく高速です。データがメモリに収まらなくなったとき、あるいはサンプルごとの変換が必要になったときにDataLoaderへ切り替えてください。

なぜ分割を層化するのか。 train_test_split(..., stratify=y)は、学習とテストで各数字の比率を同じに保ちます。テスト画像360枚で10クラスの場合、運の悪いランダム分割ではある数字のテスト例が25枚、別の数字が45枚になりかねず、クラスごとの精度がノイジーになります。

コード

学習

これをtrain.pyとして保存します。ループはエポック全体の平均学習損失と、その後のテスト精度を報告します。評価はモデルをeval()モードにしてtorch.no_grad()の下で実行します。このモデル（ドロップアウトもバッチ正規化もない）では何も変わりませんが、保つべき習慣です。

# train.py: a two-layer CNN on 16x16 digit images, trained on CPU.
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
np.random.seed(0)
device = torch.device("cpu")   # deterministic and fast enough at this size


def load_images(size: int = 16):
    """1,797 greyscale digits, 8x8 with values 0-16, upsampled to size x size in [0, 1]."""
    digits = load_digits()
    x = torch.tensor(digits.images, dtype=torch.float32).unsqueeze(1) / 16.0  # (N, 1, 8, 8)
    x = F.interpolate(x, size=(size, size), mode="bilinear", align_corners=False)
    y = torch.tensor(digits.target, dtype=torch.long)
    return x, y


class SmallCNN(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)   # 16 x 16 x 16
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)  # 32 x 8 x 8 after pool
        self.fc1 = nn.Linear(32 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, n_classes)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)   # 16 x 8 x 8
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)   # 32 x 4 x 4
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


x, y = load_images()
idx_train, idx_test = train_test_split(
    np.arange(len(y)), test_size=0.2, random_state=0, stratify=y.numpy()
)
x_train, y_train = x[idx_train].to(device), y[idx_train].to(device)
x_test, y_test = x[idx_test].to(device), y[idx_test].to(device)

model = SmallCNN().to(device)
n_params = sum(p.numel() for p in model.parameters())
optimiser = torch.optim.Adam(model.parameters(), lr=3e-3)
batch_size = 32
epochs = 5
print(f"train {len(y_train)}, test {len(y_test)}, image {tuple(x.shape[1:])}, params {n_params:,}")

t0 = time.perf_counter()
for epoch in range(1, epochs + 1):
    model.train()
    perm = torch.randperm(len(y_train))
    total_loss = 0.0
    for start in range(0, len(perm), batch_size):
        batch = perm[start:start + batch_size]
        logits = model(x_train[batch])
        loss = F.cross_entropy(logits, y_train[batch])
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
        total_loss += loss.item() * len(batch)
    model.eval()
    with torch.no_grad():
        test_acc = (model(x_test).argmax(1) == y_test).float().mean().item()
    print(f"epoch {epoch}  train loss {total_loss / len(y_train):.4f}  test acc {test_acc:.4f}")

print(f"training time {time.perf_counter() - t0:.1f} s on {device}")
torch.save(model.state_dict(), "digits_cnn.pt")
print("saved digits_cnn.pt")

期待される出力です。スクリプト全体はMac miniで3.3秒かかり、そのうち学習は0.8秒でした。残りはPyTorchのインポートです。

train 1437, test 360, image (1, 16, 16), params 38,282
epoch 1  train loss 1.7641  test acc 0.8194
epoch 2  train loss 0.3834  test acc 0.9111
epoch 3  train loss 0.2169  test acc 0.9167
epoch 4  train loss 0.1518  test acc 0.9306
epoch 5  train loss 0.0955  test acc 0.9528
training time 0.8 s on cpu
saved digits_cnn.pt

同じマシンでの2回の実行は同一の数値を出力しました。乱数のすべての源（分割、重みの初期化、バッチの順序）がシード固定され、計算がCPU上だからです。MPSやCUDAでは、小数点以下3桁目で数値がずれると考えてください。

0.511.512345学習損失テスト精度値エポック
上の出力から得たエポックごとの学習損失とテスト精度。

最初のエポックの平均損失1.76が高いのは、ネットワークが当てずっぽうの状態にある最初のバッチも含めて、エポック全体で平均しているからです（10クラスに対する一様な推測の損失はln 10、約2.30です）。エポック1の終わりには、モデルはすでにテストセットで82%に達しています。エポック5でも損失はまだ下がっており、100%との差も縮まっているので、エポックを増やせば効果があると分かります。これは意図的です。「試してみること」の節で、何エポック必要かを調べてもらいます。

混同行列で評価する

これをevaluate.pyとして保存します。モデルクラスを再度定義し、同じテスト分割を再構築し、保存した重みを読み込んで行列を出力します。digits_cnn.ptがなければ、まず同じレシピで同じモデルを学習するので、このブロックは単独で動きます。実際のプロジェクトでは、クラスは両方のスクリプトがインポートするモジュールに置きます。

# evaluate.py: reload the saved weights and print a confusion matrix.
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split

torch.manual_seed(0)


class SmallCNN(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, n_classes)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


digits = load_digits()
x = torch.tensor(digits.images, dtype=torch.float32).unsqueeze(1) / 16.0
x = F.interpolate(x, size=(16, 16), mode="bilinear", align_corners=False)
y = digits.target
_, idx_test = train_test_split(np.arange(len(y)), test_size=0.2, random_state=0, stratify=y)

model = SmallCNN()
if not Path("digits_cnn.pt").exists():
    # Same recipe as train.py, so this block also works on its own.
    idx_train, _ = train_test_split(np.arange(len(y)), test_size=0.2, random_state=0, stratify=y)
    xt, yt = x[idx_train], torch.tensor(y[idx_train])
    optimiser = torch.optim.Adam(model.parameters(), lr=3e-3)
    for epoch in range(5):
        perm = torch.randperm(len(yt))
        for start in range(0, len(perm), 32):
            b = perm[start:start + 32]
            loss = F.cross_entropy(model(xt[b]), yt[b])
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
    torch.save(model.state_dict(), "digits_cnn.pt")
model.load_state_dict(torch.load("digits_cnn.pt", map_location="cpu"))
model.eval()
with torch.no_grad():
    pred = model(x[idx_test]).argmax(1).numpy()
truth = y[idx_test]

cm = confusion_matrix(truth, pred, labels=range(10))
print("rows = true digit, columns = predicted digit")
print("     " + " ".join(f"{c:3d}" for c in range(10)))
for r in range(10):
    print(f"{r:3d}: " + " ".join(f"{v:3d}" for v in cm[r]))
per_class = cm.diagonal() / cm.sum(1)
print("per-class accuracy:", " ".join(f"{a:.2f}" for a in per_class))
print(f"overall accuracy: {(pred == truth).mean():.4f} on {len(truth)} test images")

worst = [(int(cm[r, c]), r, c) for r in range(10) for c in range(10) if r != c]
worst.sort(reverse=True)
print("most common confusions (count, true, predicted):", worst[:3])

期待される出力です。保存済みの重みがある状態でMac miniで1.7秒（先に学習が必要な場合は3.2秒）です。

rows = true digit, columns = predicted digit
       0   1   2   3   4   5   6   7   8   9
  0:  35   0   0   0   1   0   0   0   0   0
  1:   0  36   0   0   0   0   0   0   0   0
  2:   0   1  34   0   0   0   0   0   0   0
  3:   0   0   0  36   0   0   0   1   0   0
  4:   0   1   0   0  34   0   0   0   0   1
  5:   0   0   0   1   0  36   0   0   0   0
  6:   0   1   0   0   0   0  35   0   0   0
  7:   0   0   0   0   0   0   0  36   0   0
  8:   0   7   0   0   0   1   0   1  26   0
  9:   0   1   0   0   0   0   0   0   0  35
per-class accuracy: 0.97 1.00 0.97 0.97 0.94 0.97 0.97 1.00 0.74 0.97
overall accuracy: 0.9528 on 360 test images
most common confusions (count, true, predicted): [(7, 8, 1), (1, 9, 1), (1, 8, 7)]

行列は行に沿って読みます。10個の数字のうち9個は94%以上です。数字の8は74%で、35枚のテスト画像のうち7枚が1と判定されました。この1つのセルが全エラーのほぼ半分を占めています。8x8ピクセルでは細い8と1はインクの大部分を共有しており、バイリニアのアップサンプリングは輪を復元しません。1の列には2、4、6、9も1枚ずつ迷い込んでいるので、モデルには細い形状を1と予測する全般的なバイアスがあります。エポックを増やせばこの大半は直りますし（試してみてください）、最初の層を広くしても直ります。行列はどこを見るべきかを教えてくれました。精度の数値だけでは決して分からないことです。

実行

pip install torch numpy scikit-learn
python train.py
python evaluate.py

どちらのスクリプトもカレントディレクトリに書き出すので、同じ場所から実行してください。train.pyは実行のたびにdigits_cnn.ptを上書きします。

テスト

上の2つの出力がテストです。重要なチェックは次のとおりです。

  • 出力されるパラメータ数は38,282です。異なる場合は、どこかの層のサイズが間違っています。
  • 損失は毎エポック下がり、テスト精度は0.95以上で終わります。損失が約2.30で平坦なら、オプティマイザがステップしていません（よくある原因はoptimiser.step()の欠落か、zero_grad()の位置違いです）。
  • 混同行列の各行の合計は36か35で（層化分割により各数字のテスト画像は35枚か36枚）、対角成分が支配的です。
  • evaluate.pyはtrain.pyの最終エポックと同じ精度を出します。そうでなければ、2つのファイルでテスト分割か前処理がずれています。

試してみること

  • 5エポックの代わりに20エポック学習し、8対1のセルを観察してください。次に学習損失とテスト精度をプロットし、損失は下がり続けるのに精度が改善しなくなるエポックを見つけてください。その差が過学習の始まりで、嘘をつかない評価で扱っています。
  • Macではdeviceをtorch.device("mps" if torch.backends.mps.is_available() else "cpu")に、NVIDIAマシンでは"cuda"に変えて時間を計ってください。このサイズではCPUのほうがおそらく速いです。アクセラレータに処理を投入するバッチごとのオーバーヘッドが演算量を上回るからです。バッチサイズを256にしてもう一度確認してください。
  • load_digitsをtorchvision.datasets.MNIST経由のMNIST（10 MBのダウンロード）に置き換えてください。size=28とし、fc1をnn.Linear(32 * 7 * 7, 64)に変えると、CPUで数分、5エポック後に約98.5%になるはずです。
  • fc2の前にnn.Dropout(0.3)を、各畳み込みの後にnn.BatchNorm2dを追加してください。どちらもtrain()モードとeval()モードで挙動が変わる点に注意してください。ループがモードを切り替えているのはそのためです。
  • state_dict()の代わりにtorch.save(model, ...)でモデルを保存し、クラス名を変えると何が壊れるかを見てください。その後、state dictionaryの保存に戻してください。
  • 表形式データ分類器のプロジェクトのFastAPIアプリからこのモデルを提供し、16x16のリストのリストを受け付けるようにしてください。読み込みコードはlifespanフックに、順伝播はルートに置きます。

まとめ

  • プーリング付きの畳み込み2つと小さな分類器で、16x16の数字に対しCPU時間1秒以内で95%に達します。
  • パディングは畳み込みを通して空間サイズを保ちます。画像が縮むのはプーリングです。形状を追えばfc1の入力サイズは自ずと決まります。
  • 最後の桁まで再現可能な実行にしたいなら、すべてをシード固定し、デバイスを固定します。
  • state dictionaryを保存し、読み込み側のプログラムでクラスを再定義し、再読み込みしたモデルが同じテスト精度を出すことを確認します。
  • 混同行列は精度の数値1つを、具体的な問題のリストに変えます。ここでは8対1を真っ直ぐに指し示しています。