PyTorchでMNISTの手書き数字認識モデルを構築する

開発環境設定

WSL2(Ubuntu 22.04) + PyTorch 2.1.2 + Python 3.9.18環境構築

ハードウェアアクセラレーション設定


import torch

# CUDA互換GPUの利用設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用デバイス: {device}")

データ処理パイプライン構築

データセット準備


from torchvision import datasets, transforms

# 前処理パイプライン定義
transform = transforms.Compose([
    transforms.ToTensor()
])

# 訓練・評価データセットのダウンロード
train_data = datasets.MNIST(
    root='./data',
    train=True,
    transform=transform,
    download=True
)

test_data = datasets.MNIST(
    root='./data',
    train=False,
    transform=transform,
    download=True
)

データローダー構成


from torch.utils.data import DataLoader

batch_size = 64

# データローダーインスタンス生成
train_loader = DataLoader(
    dataset=train_data,
    batch_size=batch_size,
    shuffle=True
)

test_loader = DataLoader(
    dataset=test_data,
    batch_size=batch_size
)

# データ構造検証
images, labels = next(iter(train_loader))
print(f"画像テンソル形状: {images.shape}")

画像可視化実装


import matplotlib.pyplot as plt

# 20枚のサンプル画像表示
plt.figure(figsize=(15, 4))
for idx, (img, label) in enumerate(train_data[:20]):
    plt.subplot(2, 10, idx+1)
    plt.imshow(img.squeeze(), cmap='gray')
    plt.title(str(label))
    plt.axis('off')
plt.tight_layout()
plt.show()

畳み込みニューロンネットワーク設計


import torch.nn as nn

class DigitClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        # 特徴抽出層
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3),  # 畳み込み層1
            nn.ReLU(),
            nn.MaxPool2d(2),     # プーリング層1
            nn.Conv2d(16, 32, 3), # 畳み込み層2
            nn.ReLU(),
            nn.MaxPool2d(2)      # プーリング層2
        )
        
        # 分類器
        self.classifier = nn.Sequential(
            nn.Linear(32*5*5, 128), # 全結合層
            nn.ReLU(),
            nn.Linear(128, 10)      # 出力層
        )
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 32*5*5)  # テンソル平坦化
        return self.classifier(x)

# モデルインスタンス生成
model = DigitClassifier().to(device)

モデル学習プロセス

学習パラメータ設定


import torch.optim as optim

# 損失関数と最適化アルゴリズム
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

トレーニングループ実装


def train_model(model, epochs=5):
    for epoch in range(epochs):
        total_loss = 0
        for images, labels in train_loader:
            # デバイスメモリへ転送
            images, labels = images.to(device), labels.to(device)
            
            # 勾配初期化
            optimizer.zero_grad()
            
            # 順伝播計算
            outputs = model(images)
            
            # 損失計算
            loss = criterion(outputs, labels)
            
            # 逆伝播処理
            loss.backward()
            
            # パラメータ更新
            optimizer.step()
            
            total_loss += loss.item()
        
        print(f"エポック [{epoch+1}/{epochs}], 平均損失: {total_loss/len(train_loader):.4f}")

# 学習実行
train_model(model, epochs=10)

タグ: PyTorch MNIST CNN 画像分類 深層学習

6月20日 16:09 投稿