PyTorch + CNNでMNIST数字分類

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

# 1. CNNモデルの定義
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        # 畳み込み層
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)  # 28x28x1 → 28x28x32
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  # 14x14x32 → 14x14x64
        
        # プーリング層
        self.pool = nn.MaxPool2d(2, 2)  # サイズを半分に
        
        # 全結合層
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)  # 10クラス分類
        
        # ドロップアウト
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        # 畳み込み + ReLU + プーリング
        x = self.pool(torch.relu(self.conv1(x)))  # 28x28x32 → 14x14x32
        x = self.pool(torch.relu(self.conv2(x)))  # 14x14x64 → 7x7x64
        
        # 平坦化
        x = x.view(-1, 64 * 7 * 7)
        
        # 全結合層
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        
        return x

# 2. データの準備
transform = transforms.Compose([
    transforms.ToTensor(),  # PIL画像をTensorに変換
    transforms.Normalize((0.5,), (0.5,))  # 正規化 (平均, 標準偏差)
])

# データセットのダウンロードと読み込み
train_dataset = datasets.MNIST(root='~/.pytorch/data', train=True, 
                               download=True, transform=transform)
test_dataset = datasets.MNIST(root='~/.pytorch/data', train=False, 
                              download=True, transform=transform)

# DataLoader作成
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# 3. モデル、損失関数、最適化手法の設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# MPSが利用可能かチェック
if torch.backends.mps.is_available():
    device = torch.device("mps")

model = CNN().to(device)
criterion = nn.CrossEntropyLoss()  # 多クラス分類用の損失関数
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 訓練関数
def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()  # 訓練モードに設定
    total_loss = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        # 勾配をゼロに
        optimizer.zero_grad()
        
        # 順伝播
        output = model(data)
        
        # 損失計算
        loss = criterion(output, target)
        
        # 逆伝播
        loss.backward()
        
        # パラメータ更新
        optimizer.step()
        
        total_loss += loss.item()
        
        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}')
    
    return total_loss / len(train_loader)

# 5. テスト関数
def test(model, device, test_loader):
    model.eval()  # 評価モードに設定
    correct = 0
    total = 0
    
    with torch.no_grad():  # 勾配計算を無効化
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            
            # 最も確率の高いクラスを予測
            _, predicted = torch.max(output.data, 1)
            
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy

# 6. 訓練実行
epochs = 5

for epoch in range(1, epochs + 1):
    train_loss = train(model, device, train_loader, optimizer, criterion, epoch)
    print(f'Average Loss: {train_loss:.4f}')
    test(model, device, test_loader)
    print('-' * 60)

# モデルの保存
torch.save(model.state_dict(), 'mnist_cnn.pth')