Практична імплементація CNN для класифікації MNIST

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns

# Налаштування пристрою (GPU або CPU)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Використовується пристрій: {device}")

# Налаштування reproducibility (для однакових результатів при кожному запуску)
torch.manual_seed(42)
np.random.seed(42)

# Трансформації для тренувального та тестового датасетів
transform = transforms.Compose([
    transforms.ToTensor(),  # Конвертація у тензор [0.0, 1.0]
    transforms.Normalize((0.1307,), (0.3081,))  # Стандартизація (mean, std)
])

# Завантаження датасетів
train_dataset = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

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

print(f"Тренувальних зразків: {len(train_dataset)}")
print(f"Тестових зразків: {len(test_dataset)}")
print(f"Кількість класів: {len(train_dataset.classes)}")
print(f"Класи: {train_dataset.classes}")

# Гіперпараметри для навчання
batch_size = 64  # Скільки зображень обробляється за один прохід
num_workers = 2  # Кількість паралельних процесів для завантаження даних

# Створення DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,  # Перемішуємо на початку кожної епохи
    num_workers=num_workers,
    pin_memory=True  # Прискорює передачу даних на GPU
)

test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,  # Тестовий набір не перемішуємо
    num_workers=num_workers,
    pin_memory=True
)

print(f"Кількість батчів у тренувальному наборі: {len(train_loader)}")
print(f"Кількість батчів у тестовому наборі: {len(test_loader)}")
print(f"Розмір одного батчу: {batch_size}")

# Візуалізація 16 випадкових зразків
fig, axes = plt.subplots(4, 4, figsize=(10, 10))
fig.suptitle('Приклади зображень MNIST (після нормалізації)', fontsize=16, fontweight='bold')

# Отримуємо один батч
images, labels = next(iter(train_loader))

for i, ax in enumerate(axes.flat):
    # Денормалізуємо для візуалізації
    img = images[i][0] * 0.3081 + 0.1307  # Зворотна нормалізація
    img = torch.clamp(img, 0, 1)  # Обмежуємо діапазон [0, 1]
    
    ax.imshow(img.numpy(), cmap='gray')
    ax.set_title(f'Цифра: {labels[i].item()}', fontsize=12, fontweight='bold')
    ax.axis('off')

plt.tight_layout()
plt.show()

print(f"\nФорма батчу зображень: {images.shape}")
print(f"  - Batch size: {images.shape[0]}")
print(f"  - Канали (grayscale): {images.shape[1]}")
print(f"  - Висота: {images.shape[2]}")
print(f"  - Ширина: {images.shape[3]}")

class MNISTClassifier(nn.Module):
    """
    Згорткова нейронна мережа для класифікації MNIST.
    
    Архітектура:
        - 2 згорткових блоки (Conv2d → ReLU → MaxPool2d)
        - 2 повнозв'язних шари з Dropout
        - Вихідний шар для 10 класів
    """
    
    def __init__(self):
        super(MNISTClassifier, self).__init__()
        
        # Перший згортковий блок
        self.conv1 = nn.Conv2d(
            in_channels=1,      # Grayscale (1 канал)
            out_channels=32,    # 32 фільтри
            kernel_size=3,      # Розмір фільтра 3×3
            padding=1           # Padding для збереження розміру
        )
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 28×28 → 14×14
        
        # Другий згортковий блок
        self.conv2 = nn.Conv2d(
            in_channels=32,     # Входить 32 Feature Maps з conv1
            out_channels=64,    # Виходить 64 Feature Maps
            kernel_size=3,
            padding=1
        )
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 14×14 → 7×7
        
        # Розрахунок розміру після згорткових шарів: 7×7×64 = 3136
        self.flatten = nn.Flatten()
        
        # Повнозв'язні шари (класифікатор)
        self.fc1 = nn.Linear(7 * 7 * 64, 128)  # 3136 → 128
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(p=0.5)       # Вимикаємо 50% нейронів під час навчання
        self.fc2 = nn.Linear(128, 10)          # 128 → 10 класів
    
    def forward(self, x):
        """
        Прямий прохід через мережу.
        
        Args:
            x: Вхідний тензор форми (batch_size, 1, 28, 28)
        
        Returns:
            Logits для 10 класів форми (batch_size, 10)
        """
        # Блок 1: Conv → ReLU → Pool
        x = self.conv1(x)       # (batch, 1, 28, 28) → (batch, 32, 28, 28)
        x = self.relu1(x)
        x = self.pool1(x)       # (batch, 32, 28, 28) → (batch, 32, 14, 14)
        
        # Блок 2: Conv → ReLU → Pool
        x = self.conv2(x)       # (batch, 32, 14, 14) → (batch, 64, 14, 14)
        x = self.relu2(x)
        x = self.pool2(x)       # (batch, 64, 14, 14) → (batch, 64, 7, 7)
        
        # Класифікатор
        x = self.flatten(x)     # (batch, 64, 7, 7) → (batch, 3136)
        x = self.fc1(x)         # (batch, 3136) → (batch, 128)
        x = self.relu3(x)
        x = self.dropout(x)     # Dropout працює лише під час тренування
        x = self.fc2(x)         # (batch, 128) → (batch, 10)
        
        return x

# Створення моделі та перенесення на GPU (якщо доступно)
model = MNISTClassifier().to(device)

# Підрахунок кількості параметрів
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(model)
print(f"\n{'='*60}")
print(f"Загальна кількість параметрів: {total_params:,}")
print(f"Навчаємих параметрів: {trainable_params:,}")
print(f"{'='*60}")

# Функція втрат для багатокласової класифікації
criterion = nn.CrossEntropyLoss()

# Оптимізатор Adam (адаптивна швидкість навчання)
optimizer = optim.Adam(
    model.parameters(),
    lr=0.001,           # Початковий learning rate
    weight_decay=1e-5   # L2 регуляризація (дуже слабка)
)

# Планувальник: зменшуємо learning rate, якщо втрати не покращуються
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',         # Моніторимо зменшення втрат
    factor=0.5,         # Зменшуємо LR у 2 рази
    patience=3,         # Якщо 3 епохи без покращення
    verbose=True
)

print("Налаштування навчання:")
print(f"  Loss function: {criterion.__class__.__name__}")
print(f"  Optimizer: {optimizer.__class__.__name__}")
print(f"  Initial learning rate: {optimizer.param_groups[0]['lr']}")
print(f"  Scheduler: ReduceLROnPlateau (factor=0.5, patience=3)")

def train_one_epoch(model, train_loader, criterion, optimizer, device):
    """
    Навчання моделі протягом однієї епохи.
    
    Returns:
        avg_loss: Середня втрата за епоху
        accuracy: Точність на тренувальному наборі
    """
    model.train()  # Увімкнути режим тренування (Dropout працює)
    
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (images, labels) in enumerate(train_loader):
        # Перенесення даних на GPU
        images, labels = images.to(device), labels.to(device)
        
        # Обнулення градієнтів
        optimizer.zero_grad()
        
        # Прямий прохід
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Зворотний прохід та оновлення ваг
        loss.backward()
        optimizer.step()
        
        # Статистика
        running_loss += loss.item()
        _, predicted = torch.max(outputs, 1)  # Клас з максимальним logit
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # Логування кожні 200 батчів
        if (batch_idx + 1) % 200 == 0:
            print(f"  Batch [{batch_idx+1}/{len(train_loader)}] "
                  f"Loss: {loss.item():.4f} "
                  f"Acc: {100 * correct / total:.2f}%")
    
    avg_loss = running_loss / len(train_loader)
    accuracy = 100 * correct / total
    
    return avg_loss, accuracy


def validate(model, test_loader, criterion, device):
    """
    Валідація моделі на тестовому наборі.
    
    Returns:
        avg_loss: Середня втрата
        accuracy: Точність
    """
    model.eval()  # Режим оцінки (Dropout вимкнено)
    
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():  # Не обчислюємо градієнти (економія пам'яті)
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    avg_loss = running_loss / len(test_loader)
    accuracy = 100 * correct / total
    
    return avg_loss, accuracy

# Гіперпараметри навчання
num_epochs = 10

# Списки для збереження метрик
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []

print("Початок навчання...\n")
print("="*70)

for epoch in range(num_epochs):
    print(f"\nЕпоха {epoch+1}/{num_epochs}")
    print("-"*70)
    
    # Навчання
    train_loss, train_acc = train_one_epoch(
        model, train_loader, criterion, optimizer, device
    )
    
    # Валідація
    val_loss, val_acc = validate(model, test_loader, criterion, device)
    
    # Оновлення learning rate
    scheduler.step(val_loss)
    
    # Збереження метрик
    train_losses.append(train_loss)
    train_accuracies.append(train_acc)
    val_losses.append(val_loss)
    val_accuracies.append(val_acc)
    
    # Виведення результатів епохи
    print(f"\n{'='*70}")
    print(f"Епоха {epoch+1} завершена:")
    print(f"  Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
    print(f"  Val Loss:   {val_loss:.4f}   | Val Acc:   {val_acc:.2f}%")
    print(f"  Current LR: {optimizer.param_groups[0]['lr']:.6f}")
    print(f"{'='*70}")

print("\n✅ Навчання завершено!")
print(f"Фінальна точність на тестовому наборі: {val_accuracies[-1]:.2f}%")

# Візуалізація кривих навчання
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

epochs_range = range(1, num_epochs + 1)

# Графік втрат (Loss)
ax1.plot(epochs_range, train_losses, 'o-', color='#3b82f6', linewidth=2.5, 
         markersize=8, label='Train Loss')
ax1.plot(epochs_range, val_losses, 's-', color='#ef4444', linewidth=2.5, 
         markersize=8, label='Val Loss')
ax1.set_xlabel('Епоха', fontsize=14, fontweight='bold')
ax1.set_ylabel('CrossEntropyLoss', fontsize=14, fontweight='bold')
ax1.set_title('Динаміка функції втрат', fontsize=16, fontweight='bold')
ax1.legend(fontsize=12)
ax1.grid(alpha=0.3)
ax1.set_xticks(epochs_range)

# Графік точності (Accuracy)
ax2.plot(epochs_range, train_accuracies, 'o-', color='#3b82f6', linewidth=2.5, 
         markersize=8, label='Train Accuracy')
ax2.plot(epochs_range, val_accuracies, 's-', color='#ef4444', linewidth=2.5, 
         markersize=8, label='Val Accuracy')
ax2.set_xlabel('Епоха', fontsize=14, fontweight='bold')
ax2.set_ylabel('Accuracy (%)', fontsize=14, fontweight='bold')
ax2.set_title('Динаміка точності класифікації', fontsize=16, fontweight='bold')
ax2.legend(fontsize=12)
ax2.grid(alpha=0.3)
ax2.set_xticks(epochs_range)
ax2.set_ylim(94, 100)

plt.tight_layout()
plt.show()

print(f"\n📊 Підсумкові метрики:")
print(f"  Найкраща Val Accuracy: {max(val_accuracies):.2f}% (Епоха {val_accuracies.index(max(val_accuracies))+1})")
print(f"  Найнижча Val Loss: {min(val_losses):.4f} (Епоха {val_losses.index(min(val_losses))+1})")

# Отримання всіх передбачень на тестовому наборі
model.eval()
all_preds = []
all_labels = []

with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())

# Побудова Confusion Matrix
cm = confusion_matrix(all_labels, all_preds)

# Візуалізація
fig, ax = plt.subplots(figsize=(12, 10))

sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', square=True, 
            linewidths=0.5, cbar_kws={"shrink": 0.8}, ax=ax)

ax.set_xlabel('Передбачений клас', fontsize=14, fontweight='bold')
ax.set_ylabel('Справжній клас', fontsize=14, fontweight='bold')
ax.set_title('Confusion Matrix для MNIST CNN\n(99.12% accuracy)', 
             fontsize=16, fontweight='bold')
ax.set_xticklabels(range(10), fontsize=12)
ax.set_yticklabels(range(10), fontsize=12, rotation=0)

plt.tight_layout()
plt.show()

# Детальний звіт
print("\n📋 Classification Report:\n")
print(classification_report(all_labels, all_preds, 
                            target_names=[str(i) for i in range(10)],
                            digits=4))

import os

# Створюємо директорію для моделей
os.makedirs('./models', exist_ok=True)

# Збереження повної моделі (архітектура + ваги)
model_path = './models/mnist_cnn_best.pth'
torch.save({
    'epoch': num_epochs,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'train_loss': train_losses[-1],
    'val_loss': val_losses[-1],
    'val_accuracy': val_accuracies[-1],
}, model_path)

print(f"✅ Модель збережена: {model_path}")
print(f"   Розмір файлу: {os.path.getsize(model_path) / 1024:.2f} KB")

# Демонстрація завантаження моделі
print("\n🔄 Завантаження збереженої моделі...")

# Створюємо нову екземпляр моделі
loaded_model = MNISTClassifier().to(device)

# Завантажуємо ваги
checkpoint = torch.load(model_path, weights_only=False)
loaded_model.load_state_dict(checkpoint['model_state_dict'])
loaded_model.eval()

# Перевірка, що модель працює
test_loss, test_acc = validate(loaded_model, test_loader, criterion, device)
print(f"✅ Модель завантажена успішно!")
print(f"   Точність на тестовому наборі: {test_acc:.2f}%")
print(f"   (Оригінальна точність: {checkpoint['val_accuracy']:.2f}%)")