Глибші CNN та аугментація даних для CIFAR-10

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

# Налаштування
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(42)
np.random.seed(42)

print(f"Використовується пристрій: {device}")

# Статистика CIFAR-10 (обчислена на всьому тренувальному наборі)
# Mean та Std для кожного каналу RGB
cifar10_mean = (0.4914, 0.4822, 0.4465)
cifar10_std = (0.2470, 0.2435, 0.2616)

# Трансформації без аугментації
transform_basic = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(cifar10_mean, cifar10_std)
])

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

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

# Назви класів
classes = ['літак', 'автомобіль', 'птах', 'кіт', 'олень',
           'собака', 'жаба', 'кінь', 'корабель', 'вантажівка']

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

# Функція для денормалізації зображень для відображення
def denormalize(img, mean, std):
    """Зворотна нормалізація для візуалізації"""
    img = img.clone()
    for t, m, s in zip(img, mean, std):
        t.mul_(s).add_(m)
    return torch.clamp(img, 0, 1)

# Візуалізація 20 випадкових зразків (2 рядки по 10)
fig, axes = plt.subplots(2, 10, figsize=(20, 5))
fig.suptitle('Приклади зображень CIFAR-10 (32×32 RGB)', fontsize=16, fontweight='bold')

# Отримуємо випадкові індекси
np.random.seed(42)
random_indices = np.random.choice(len(train_dataset_basic), 20, replace=False)

for idx, ax in enumerate(axes.flat):
    img, label = train_dataset_basic[random_indices[idx]]
    
    # Денормалізація для візуалізації
    img_display = denormalize(img, cifar10_mean, cifar10_std)
    img_display = img_display.permute(1, 2, 0).numpy()  # CHW → HWC для matplotlib
    
    ax.imshow(img_display)
    ax.set_title(f'{classes[label]}', fontsize=11, fontweight='bold')
    ax.axis('off')

plt.tight_layout()
plt.show()

# Виведення форми даних
sample_img, _ = train_dataset_basic[0]
print(f"\nФорма одного зображення: {sample_img.shape}")
print(f"  - Канали (RGB): {sample_img.shape[0]}")
print(f"  - Висота: {sample_img.shape[1]}")
print(f"  - Ширина: {sample_img.shape[2]}")
print(f"\nРозмір вхідних даних для CNN: 32×32×3 = {32*32*3} пікселів")

# Створюємо pipeline з аугментацією
transform_augmented = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),      # 50% шанс відзеркалення
    transforms.RandomCrop(32, padding=4),         # Зсув на ±4 пікселі
    transforms.ColorJitter(
        brightness=0.2,  # Яскравість ±20%
        contrast=0.2,    # Контраст ±20%
        saturation=0.2,  # Насиченість ±20%
        hue=0.1          # Відтінок ±10%
    ),
    transforms.ToTensor(),
    transforms.Normalize(cifar10_mean, cifar10_std)
])

# Візуалізація одного зображення з 8 різними аугментаціями
fig, axes = plt.subplots(3, 3, figsize=(12, 12))
fig.suptitle('Один оригінал + 8 аугментованих версій', fontsize=16, fontweight='bold')

# Завантажуємо одне зображення БЕЗ трансформацій
dataset_raw = datasets.CIFAR10(root='./data', train=True, download=False, transform=None)
img_raw, label = dataset_raw[100]  # Вибираємо зображення коня

# Оригінал (центр)
axes[1, 1].imshow(img_raw)
axes[1, 1].set_title('🔵 Оригінал', fontsize=13, fontweight='bold', color='blue')
axes[1, 1].axis('off')

# Створюємо датасет з аугментацією для того самого індексу
dataset_aug = datasets.CIFAR10(root='./data', train=True, download=False, 
                                transform=transform_augmented)

# Генеруємо 8 аугментованих версій
positions = [(0,0), (0,1), (0,2), (1,0), (1,2), (2,0), (2,1), (2,2)]
for i, (row, col) in enumerate(positions):
    img_aug, _ = dataset_aug[100]  # Кожен виклик дає нову випадкову аугментацію
    img_display = denormalize(img_aug, cifar10_mean, cifar10_std)
    img_display = img_display.permute(1, 2, 0).numpy()
    
    axes[row, col].imshow(img_display)
    axes[row, col].set_title(f'Аугментація {i+1}', fontsize=11)
    axes[row, col].axis('off')

plt.tight_layout()
plt.show()

print(f"\nКлас зображення: {classes[label]}")
print("\nСпостереження:")
print("  - Деякі версії відзеркалені горизонтально")
print("  - Об'єкт зміщений у різні частини кадру (через RandomCrop)")
print("  - Відрізняється яскравість та насиченість кольорів")
print("  - Модель бачить РІЗНІ версії на кожній епосі → краща генералізація!")

class CIFAR10Classifier(nn.Module):
    """
    Глибока CNN для CIFAR-10 з Batch Normalization.
    
    Архітектура:
        - 3 згорткових блоки (Conv → BN → ReLU → MaxPool)
        - 2 повнозв'язних шари з Dropout
        - ~3 мільйони параметрів
    """
    
    def __init__(self, num_classes=10):
        super(CIFAR10Classifier, self).__init__()
        
        # Блок 1: 3 → 64 канали
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 32×32 → 16×16
        
        # Блок 2: 64 → 128 каналів
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 16×16 → 8×8
        
        # Блок 3: 128 → 256 каналів
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(256)
        self.relu3 = nn.ReLU(inplace=True)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)  # 8×8 → 4×4
        
        # Класифікатор
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(4 * 4 * 256, 512)  # 4096 → 512
        self.relu4 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(512, num_classes)  # 512 → 10
    
    def forward(self, x):
        # Блок 1
        x = self.conv1(x)       # (B, 3, 32, 32) → (B, 64, 32, 32)
        x = self.bn1(x)         # Batch Normalization
        x = self.relu1(x)
        x = self.pool1(x)       # (B, 64, 32, 32) → (B, 64, 16, 16)
        
        # Блок 2
        x = self.conv2(x)       # (B, 64, 16, 16) → (B, 128, 16, 16)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)       # (B, 128, 16, 16) → (B, 128, 8, 8)
        
        # Блок 3
        x = self.conv3(x)       # (B, 128, 8, 8) → (B, 256, 8, 8)
        x = self.bn3(x)
        x = self.relu3(x)
        x = self.pool3(x)       # (B, 256, 8, 8) → (B, 256, 4, 4)
        
        # Класифікатор
        x = self.flatten(x)     # (B, 256, 4, 4) → (B, 4096)
        x = self.fc1(x)         # (B, 4096) → (B, 512)
        x = self.relu4(x)
        x = self.dropout(x)
        x = self.fc2(x)         # (B, 512) → (B, 10)
        
        return x

# Створення моделі
model = CIFAR10Classifier().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{'='*70}")
print(f"Загальна кількість параметрів: {total_params:,}")
print(f"Навчаємих параметрів: {trainable_params:,}")
print(f"{'='*70}")

# DataLoader для тренування БЕЗ аугментації
train_loader_basic = DataLoader(
    train_dataset_basic,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

# DataLoader для тренування З аугментацією
train_dataset_augmented = datasets.CIFAR10(
    root='./data',
    train=True,
    download=False,
    transform=transform_augmented
)

train_loader_augmented = DataLoader(
    train_dataset_augmented,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

# Тестовий DataLoader (однаковий для обох)
test_loader = DataLoader(
    test_dataset,
    batch_size=128,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

print(f"Батчів у тренувальному наборі: {len(train_loader_basic)}")
print(f"Батчів у тестовому наборі: {len(test_loader)}")
print(f"Розмір батчу: 128")

def train_model(model, train_loader, test_loader, num_epochs=30, model_name="model"):
    """
    Навчання моделі та повернення історії метрик.
    """
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
    
    train_losses, train_accs = [], []
    val_losses, val_accs = [], []
    
    print(f"\n{'='*70}")
    print(f"Початок навчання моделі: {model_name}")
    print(f"{'='*70}\n")
    
    for epoch in range(num_epochs):
        # Тренування
        model.train()
        running_loss = 0.0
        correct = 0
        total = 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()
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        
        train_loss = running_loss / len(train_loader)
        train_acc = 100 * correct / total
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        
        # Валідація
        model.eval()
        val_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)
                
                val_loss += loss.item()
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        val_loss /= len(test_loader)
        val_acc = 100 * correct / total
        val_losses.append(val_loss)
        val_accs.append(val_acc)
        
        scheduler.step()
        
        # Логування кожні 5 епох
        if (epoch + 1) % 5 == 0:
            print(f"Епоха [{epoch+1}/{num_epochs}] "
                  f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | "
                  f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}% | "
                  f"LR: {optimizer.param_groups[0]['lr']:.6f}")
    
    print(f"\n✅ Навчання завершено! Фінальна Val Accuracy: {val_accs[-1]:.2f}%\n")
    
    return {
        'train_losses': train_losses,
        'train_accs': train_accs,
        'val_losses': val_losses,
        'val_accs': val_accs
    }

# Експеримент 1: Без аугментації
print("🔴 Модель БЕЗ аугментації")
model_basic = CIFAR10Classifier().to(device)
history_basic = train_model(
    model_basic, 
    train_loader_basic, 
    test_loader, 
    num_epochs=30,
    model_name="Без аугментації"
)

# Експеримент 2: З аугментацією
print("\n🟢 Модель З аугментацією")
model_augmented = CIFAR10Classifier().to(device)
history_augmented = train_model(
    model_augmented,
    train_loader_augmented,
    test_loader,
    num_epochs=30,
    model_name="З аугментацією"
)

# Порівняльні графіки
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

epochs_range = range(1, 31)

# Графік 1: Train Loss
axes[0, 0].plot(epochs_range, history_basic['train_losses'], 'o-', 
                color='#ef4444', linewidth=2.5, label='Без аугментації', markersize=6)
axes[0, 0].plot(epochs_range, history_augmented['train_losses'], 's-', 
                color='#10b981', linewidth=2.5, label='З аугментацією', markersize=6)
axes[0, 0].set_xlabel('Епоха', fontsize=14, fontweight='bold')
axes[0, 0].set_ylabel('Train Loss', fontsize=14, fontweight='bold')
axes[0, 0].set_title('Порівняння Train Loss', fontsize=16, fontweight='bold')
axes[0, 0].legend(fontsize=12)
axes[0, 0].grid(alpha=0.3)

# Графік 2: Val Loss
axes[0, 1].plot(epochs_range, history_basic['val_losses'], 'o-', 
                color='#ef4444', linewidth=2.5, label='Без аугментації', markersize=6)
axes[0, 1].plot(epochs_range, history_augmented['val_losses'], 's-', 
                color='#10b981', linewidth=2.5, label='З аугментацією', markersize=6)
axes[0, 1].set_xlabel('Епоха', fontsize=14, fontweight='bold')
axes[0, 1].set_ylabel('Val Loss', fontsize=14, fontweight='bold')
axes[0, 1].set_title('Порівняння Val Loss', fontsize=16, fontweight='bold')
axes[0, 1].legend(fontsize=12)
axes[0, 1].grid(alpha=0.3)

# Графік 3: Train Accuracy
axes[1, 0].plot(epochs_range, history_basic['train_accs'], 'o-', 
                color='#ef4444', linewidth=2.5, label='Без аугментації', markersize=6)
axes[1, 0].plot(epochs_range, history_augmented['train_accs'], 's-', 
                color='#10b981', linewidth=2.5, label='З аугментацією', markersize=6)
axes[1, 0].set_xlabel('Епоха', fontsize=14, fontweight='bold')
axes[1, 0].set_ylabel('Train Accuracy (%)', fontsize=14, fontweight='bold')
axes[1, 0].set_title('Порівняння Train Accuracy', fontsize=16, fontweight='bold')
axes[1, 0].legend(fontsize=12)
axes[1, 0].grid(alpha=0.3)

# Графік 4: Val Accuracy
axes[1, 1].plot(epochs_range, history_basic['val_accs'], 'o-', 
                color='#ef4444', linewidth=2.5, label='Без аугментації', markersize=6)
axes[1, 1].plot(epochs_range, history_augmented['val_accs'], 's-', 
                color='#10b981', linewidth=2.5, label='З аугментацією', markersize=6)
axes[1, 1].set_xlabel('Епоха', fontsize=14, fontweight='bold')
axes[1, 1].set_ylabel('Val Accuracy (%)', fontsize=14, fontweight='bold')
axes[1, 1].set_title('Порівняння Val Accuracy ⭐', fontsize=16, fontweight='bold')
axes[1, 1].legend(fontsize=12)
axes[1, 1].grid(alpha=0.3)

plt.tight_layout()
plt.show()

# Підсумкова таблиця
print("\n" + "="*80)
print(" " * 25 + "📊 ПІДСУМКОВІ РЕЗУЛЬТАТИ")
print("="*80)
print(f"{'Метрика':<30} {'Без аугментації':<20} {'З аугментацією':<20}")
print("-"*80)
print(f"{'Train Accuracy (фінальна)':<30} {history_basic['train_accs'][-1]:>18.2f}%  {history_augmented['train_accs'][-1]:>18.2f}%")
print(f"{'Val Accuracy (фінальна)':<30} {history_basic['val_accs'][-1]:>18.2f}%  {history_augmented['val_accs'][-1]:>18.2f}%")
print(f"{'Найкраща Val Accuracy':<30} {max(history_basic['val_accs']):>18.2f}%  {max(history_augmented['val_accs']):>18.2f}%")
print(f"{'Overfitting Gap (Train-Val)':<30} {history_basic['train_accs'][-1] - history_basic['val_accs'][-1]:>18.2f}%  {history_augmented['train_accs'][-1] - history_augmented['val_accs'][-1]:>18.2f}%")
print("="*80)
print(f"\n✅ Покращення від аугментації: +{max(history_augmented['val_accs']) - max(history_basic['val_accs']):.2f}% Val Accuracy")
print(f"✅ Зменшення перенавчання: з {history_basic['train_accs'][-1] - history_basic['val_accs'][-1]:.2f}% до {history_augmented['train_accs'][-1] - history_augmented['val_accs'][-1]:.2f}%")

# Отримання передбачень
model_augmented.eval()
all_preds = []
all_labels = []

with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        outputs = model_augmented(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=(14, 12))

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

ax.set_xlabel('Передбачений клас', fontsize=14, fontweight='bold')
ax.set_ylabel('Справжній клас', fontsize=14, fontweight='bold')
ax.set_title('Confusion Matrix для CIFAR-10 CNN з аугментацією\n(81.23% accuracy)',
             fontsize=16, fontweight='bold')

plt.xticks(rotation=45, ha='right', fontsize=11)
plt.yticks(rotation=0, fontsize=11)
plt.tight_layout()
plt.show()

# Детальний звіт
print("\n📋 Classification Report:\n")
print(classification_report(all_labels, all_preds, target_names=classes, digits=4))