951 views

この記事は最終更新から 397日 が経過しています。

【1】やりたいこと

過去記事で使用した PyTorch公式の MNISTサンプルプログラムは基礎学習には少々大きめだ。
必要機能だけを切り出して、ミニマム実装してみたい。

プログラムの最小化により、
処理のエッセンスが明確にわかるようになるだろう。

【2】やってみた

(1) CPU実行用の実装

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

#//////////////////////////////////////////////////////////////////////////////
# 1. データの準備
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

dataDir = '../data'
train_dataset = datasets.MNIST(dataDir, train=True, download=True, transform=transform)
test_dataset  = datasets.MNIST(dataDir, train=False,               transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64,   shuffle=True)
test_loader  = DataLoader(test_dataset,  batch_size=1000, shuffle=False)

#//////////////////////////////////////////////////////////////////////////////
# 2. ニューラルネットワークの定義
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)      # [BatchSize][28][28] → [BatchSize][784] 1次元ベクトルに変換
        x = torch.relu(self.fc1(x))  #  → [BatchSize][128] → ReLU
        x = torch.relu(self.fc2(x))  #  → [BatchSize][64] → ReLU
        x = self.fc3(x)              #  → [BatchSize][10] この後、criterionで softmax処理される。
        return x

model = Net()

#//////////////////////////////////////////////////////////////////////////////
# 3. 損失関数・最適化手法の定義
criterion = nn.CrossEntropyLoss()    # 分類タスク向けの損失関数
optimizer = optim.Adam(model.parameters())

#//////////////////////////////////////////////////////////////////////////////
# 4. 学習ループ
for epoch in range(5):
    model.train()                          # 学習モードに設定 (dropout等を許可)
    for data, target in train_loader:      # イテレータからデータを取り出す。
        optimizer.zero_grad()              # 勾配をゼロクリア
        output = model(data)               # 順伝播を実行
        loss = criterion(output, target)   # 損失関数を実行 (誤差を算出)
        loss.backward()                    # 逆伝播
        optimizer.step()                   # 各パラメータの更新
    print(f"Epoch {epoch+1} complete")

#//////////////////////////////////////////////////////////////////////////////
# 5. テスト
model.eval()                               # テストモードに設定 (dropout等を禁止)
correct = 0                                # 正解数 = 0
total = 0                                  # テスト数 = 0
with torch.no_grad():                      # テスト中は勾配を計算しないようにする。（メモリ節約）
    for data, target in test_loader:       # イテレータからデータを取り出す。
        outputs = model(data)              # 順伝播
        _, predicted = torch.max(outputs.data, 1)      # 出力層の最大値のインデックス（=ラベル番号）を取得
        total += target.size(0)                        # テスト数を更新
        correct += (predicted == target).sum().item()  # 正解数を更新（比較結果がTrueの数を加算）
print(f"Test Accuracy: {100 * correct / total:.2f}%")

処理のポイントについて、メモ書きしておく。

9-12行目：
入力データの前処理を定義する。
画像をテンソルに変換し、平均 0.1307、標準偏差 0.3081 で正規化している。
これにより学習が収束しやすくなり、安定化・高速化が期待できる。

15-18行目：
MNISTデータセットをファイルからロードし、イテレータ型オブジェクトに格納している。
指定ディレクトリにデータセットファイルが存在しなければ、自動的にダウンロードしてくる。
イテレータからデータを取得時に、上記の正規化済みのデータが取得できる。
イテレータからデータを取得時に、学習データはシャッフルした後、テストデータはシャッフルせずに取得する。

実行結果は以下の通り。

$ python mnist_samnple_CPU.py
Epoch 1 complete
Epoch 2 complete
Epoch 3 complete
Epoch 4 complete
Epoch 5 complete
Test Accuracy: 97.48%

(2) GPU実行用の実装

変更箇所をハイライト表示する。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

#//////////////////////////////////////////////////////////////////////////////
# 1. データの準備
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

dataDir = '../data'
train_dataset = datasets.MNIST(dataDir, train=True, download=True, transform=transform)
test_dataset  = datasets.MNIST(dataDir, train=False,               transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64,   shuffle=True)
test_loader  = DataLoader(test_dataset,  batch_size=1000, shuffle=False)

#//////////////////////////////////////////////////////////////////////////////
# 2. ニューラルネットワークの定義
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
model = Net().to(device)

#//////////////////////////////////////////////////////////////////////////////
# 3. 損失関数・最適化手法の定義
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

#//////////////////////////////////////////////////////////////////////////////
# 4. 学習ループ
for epoch in range(5):  # エポック数はお好みで調整
    model.train()
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1} complete")

#//////////////////////////////////////////////////////////////////////////////
# 5. テスト
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        outputs = model(data)
        _, predicted = torch.max(outputs.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()
print(f"Test Accuracy: {100 * correct / total:.2f}%")

実行結果は以下の通り。
CPU版と比較するとだいぶ高速だ。

$ python mnist_samnple_GPU.py
Device: cuda
Epoch 1 complete
Epoch 2 complete
Epoch 3 complete
Epoch 4 complete
Epoch 5 complete
Test Accuracy: 96.98%