Batch Normalization — прорив у глибокому навчанні

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

# Проста мережа БЕЗ Batch Normalization
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 256)  # Будемо моніторити цей шар
        self.fc4 = nn.Linear(256, 128)
        self.fc5 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x, return_layer3=False):
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        layer3_input = x.clone()  # Зберегти вхід 3-го шару
        x = self.relu(self.fc3(x))
        x = self.relu(self.fc4(x))
        x = self.fc5(x)
        
        if return_layer3:
            return x, layer3_input
        return x

# Ініціалізація
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# MNIST DataLoader
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

# Збереження статистики розподілу входів 3-го шару
layer3_means = []
layer3_stds = []

model.train()
for iteration, (images, labels) in enumerate(train_loader):
    if iteration >= 100:  # Перші 100 ітерацій
        break
    
    images, labels = images.to(device), labels.to(device)
    
    # Forward з збереженням активацій
    outputs, layer3_input = model(images, return_layer3=True)
    loss = criterion(outputs, labels)
    
    # Backward та оптимізація
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Статистика входу 3-го шару
    with torch.no_grad():
        mean = layer3_input.mean().item()
        std = layer3_input.std().item()
        layer3_means.append(mean)
        layer3_stds.append(std)

print(f"✅ Відстежено {len(layer3_means)} ітерацій навчання")
print(f"   Початкове середнє входу 3-го шару: {layer3_means[0]:.4f}")
print(f"   Фінальне середнє входу 3-го шару:  {layer3_means[-1]:.4f}")
print(f"   Зміна: {abs(layer3_means[-1] - layer3_means[0]):.4f}")

# Візуалізація дрейфу розподілу
fig, axes = plt.subplots(1, 2, figsize=(16, 5))

# Графік 1: Зміна середнього
ax1 = axes[0]
ax1.plot(range(len(layer3_means)), layer3_means, 
         color='#3b82f6', linewidth=2, label='Mean (μ)')
ax1.axhline(layer3_means[0], color='gray', linestyle='--', 
            alpha=0.5, label=f'Початкове μ = {layer3_means[0]:.3f}')
ax1.fill_between(range(len(layer3_means)), 
                 layer3_means[0], layer3_means,
                 alpha=0.2, color='#3b82f6')
ax1.set_xlabel('Ітерація', fontsize=12)
ax1.set_ylabel('Середнє значення входу', fontsize=12)
ax1.set_title('Дрейф середнього (μ) протягом навчання', 
              fontsize=14, fontweight='bold')
ax1.legend(fontsize=11)
ax1.grid(alpha=0.3)

# Графік 2: Зміна стандартного відхилення
ax2 = axes[1]
ax2.plot(range(len(layer3_stds)), layer3_stds, 
         color='#ef4444', linewidth=2, label='Std (σ)')
ax2.axhline(layer3_stds[0], color='gray', linestyle='--', 
            alpha=0.5, label=f'Початкове σ = {layer3_stds[0]:.3f}')
ax2.fill_between(range(len(layer3_stds)), 
                 layer3_stds[0], layer3_stds,
                 alpha=0.2, color='#ef4444')
ax2.set_xlabel('Ітерація', fontsize=12)
ax2.set_ylabel('Стандартне відхилення входу', fontsize=12)
ax2.set_title('Дрейф дисперсії (σ²) протягом навчання', 
              fontsize=14, fontweight='bold')
ax2.legend(fontsize=11)
ax2.grid(alpha=0.3)

plt.tight_layout()
plt.show()

import torch
import torch.nn as nn

# Простий шар з BN
layer = nn.Sequential(
    nn.Linear(10, 10),
    nn.BatchNorm1d(10),
    nn.ReLU()
)

# Створимо один приклад (без батчу)
x = torch.randn(1, 10)

print("🔬 Експеримент: один і той же вхід у різних режимах\n")

# Train mode: кожен forward pass дає різні результати!
layer.train()
out1_train = layer(x)
out2_train = layer(x)
out3_train = layer(x)

print("📊 Training Mode:")
print(f"   1-й forward: {out1_train[0, :3].detach().numpy()}")
print(f"   2-й forward: {out2_train[0, :3].detach().numpy()}")
print(f"   3-й forward: {out3_train[0, :3].detach().numpy()}")
print(f"   Різниця 1-2: {torch.norm(out1_train - out2_train).item():.6f}")
print(f"   Різниця 1-3: {torch.norm(out1_train - out3_train).item():.6f}\n")

# Eval mode: forward pass детермінований
layer.eval()
out1_eval = layer(x)
out2_eval = layer(x)
out3_eval = layer(x)

print("📊 Evaluation Mode:")
print(f"   1-й forward: {out1_eval[0, :3].detach().numpy()}")
print(f"   2-й forward: {out2_eval[0, :3].detach().numpy()}")
print(f"   3-й forward: {out3_eval[0, :3].detach().numpy()}")
print(f"   Різниця 1-2: {torch.norm(out1_eval - out2_eval).item():.6f}")
print(f"   Різниця 1-3: {torch.norm(out1_eval - out3_eval).item():.6f}")

class MLPWithBN(nn.Module):
    def __init__(self):
        super(MLPWithBN, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.bn1 = nn.BatchNorm1d(256)  # BN для 256 features
        
        self.fc2 = nn.Linear(256, 128)
        self.bn2 = nn.BatchNorm1d(128)
        
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten
        
        x = self.fc1(x)
        x = self.bn1(x)  # Нормалізація
        x = self.relu(x)
        
        x = self.fc2(x)
        x = self.bn2(x)
        x = self.relu(x)
        
        x = self.fc3(x)  # Останній шар БЕЗ BN
        return x

model = MLPWithBN()
print(model)
print(f"\n📊 Параметри BatchNorm1d:")
print(f"   bn1.weight (γ): {model.bn1.weight.shape}")  # (256,)
print(f"   bn1.bias (β):   {model.bn1.bias.shape}")    # (256,)
print(f"   bn1.running_mean: {model.bn1.running_mean.shape}")  # (256,)
print(f"   bn1.running_var:  {model.bn1.running_var.shape}")   # (256,)

class CNNWithBN(nn.Module):
    def __init__(self):
        super(CNNWithBN, self).__init__()
        # Блок 1
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)  # BN для 64 каналів
        
        # Блок 2
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        
        # Класифікатор
        self.pool = nn.MaxPool2d(2, 2)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(128 * 8 * 8, 10)
    
    def forward(self, x):
        # Блок 1: Conv → BN → ReLU → Pool
        x = self.conv1(x)      # (batch, 64, 32, 32)
        x = self.bn1(x)        # Нормалізація по каналах
        x = self.relu(x)
        x = self.pool(x)       # (batch, 64, 16, 16)
        
        # Блок 2
        x = self.conv2(x)      # (batch, 128, 16, 16)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.pool(x)       # (batch, 128, 8, 8)
        
        # Класифікатор
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = CNNWithBN()
print(f"📊 Параметри BatchNorm2d:")
print(f"   bn1.weight (γ): {model.bn1.weight.shape}")  # (64,) — по одному на канал
print(f"   bn1.running_mean: {model.bn1.running_mean.shape}")  # (64,)

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

# Модель БЕЗ Batch Normalization
class CNNBaseline(nn.Module):
    def __init__(self):
        super(CNNBaseline, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)  # Потрібен Dropout для регуляризації
    
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        x = self.relu(self.conv2(x))
        x = self.pool(x)
        x = self.relu(self.conv3(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.fc2(x)
        return x

# Модель З Batch Normalization
class CNNWithBN(nn.Module):
    def __init__(self):
        super(CNNWithBN, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        
        self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(256)
        
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.bn_fc = nn.BatchNorm1d(512)
        self.fc2 = nn.Linear(512, 10)
        self.relu = nn.ReLU()
        # Dropout менший або взагалі не потрібен з BN
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.pool(x)
        
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu(x)
        x = self.pool(x)
        
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.bn_fc(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# Завантаження CIFAR-10
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

train_dataset = datasets.CIFAR10('./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10('./data', train=False, download=True, transform=transform_test)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

print("✅ Датасет завантажено")
print(f"   Train: {len(train_dataset)} зображень")
print(f"   Test:  {len(test_dataset)} зображень")

def train_and_evaluate(model, model_name, train_loader, test_loader, 
                      epochs=20, lr=0.01, device='cuda'):
    """Навчання моделі та збереження метрик"""
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    
    history = {
        'train_loss': [], 'train_acc': [],
        'test_loss': [], 'test_acc': [],
        'epoch_time': []
    }
    
    print(f"\n🚀 Навчання моделі: {model_name}")
    print(f"   Learning rate: {lr}")
    
    for epoch in range(epochs):
        start_time = time.time()
        
        # Training
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            train_total += labels.size(0)
            train_correct += predicted.eq(labels).sum().item()
        
        # Evaluation
        model.eval()
        test_loss = 0.0
        test_correct = 0
        test_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)
                
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                test_total += labels.size(0)
                test_correct += predicted.eq(labels).sum().item()
        
        epoch_time = time.time() - start_time
        train_loss_avg = train_loss / len(train_loader)
        train_acc = 100. * train_correct / train_total
        test_loss_avg = test_loss / len(test_loader)
        test_acc = 100. * test_correct / test_total
        
        history['train_loss'].append(train_loss_avg)
        history['train_acc'].append(train_acc)
        history['test_loss'].append(test_loss_avg)
        history['test_acc'].append(test_acc)
        history['epoch_time'].append(epoch_time)
        
        scheduler.step()
        
        if (epoch + 1) % 5 == 0:
            print(f"   Епоха {epoch+1:2d}/{epochs}: "
                  f"Train Acc={train_acc:.2f}%, Test Acc={test_acc:.2f}%, "
                  f"Time={epoch_time:.1f}s")
    
    return history

# Експеримент
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Модель 1: Baseline без BN
torch.manual_seed(42)
model_baseline = CNNBaseline()
history_baseline = train_and_evaluate(
    model_baseline, "Baseline (без BN)", 
    train_loader, test_loader, 
    epochs=20, lr=0.01, device=device
)

# Модель 2: З Batch Normalization
torch.manual_seed(42)
model_bn = CNNWithBN()
history_bn = train_and_evaluate(
    model_bn, "CNN + BatchNorm", 
    train_loader, test_loader, 
    epochs=20, lr=0.1,  # ← ЗБІЛЬШЕНИЙ LR у 10 разів!
    device=device
)

print("\n🎉 Експеримент завершено!")

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Графік 1: Test Accuracy
ax1 = axes[0]
ax1.plot(range(1, 21), history_baseline['test_acc'], 
         marker='o', linewidth=2.5, label='Baseline (LR=0.01)', color='#ef4444')
ax1.plot(range(1, 21), history_bn['test_acc'], 
         marker='s', linewidth=2.5, label='+ BatchNorm (LR=0.1)', color='#22c55e')
ax1.set_xlabel('Епоха', fontsize=13)
ax1.set_ylabel('Test Accuracy (%)', fontsize=13)
ax1.set_title('Точність на тестовому датасеті', fontsize=15, fontweight='bold')
ax1.legend(fontsize=12)
ax1.grid(alpha=0.3)
ax1.set_ylim([40, 90])

# Графік 2: Train Loss
ax2 = axes[1]
ax2.plot(range(1, 21), history_baseline['train_loss'], 
         marker='o', linewidth=2.5, label='Baseline', color='#ef4444')
ax2.plot(range(1, 21), history_bn['train_loss'], 
         marker='s', linewidth=2.5, label='+ BatchNorm', color='#22c55e')
ax2.set_xlabel('Епоха', fontsize=13)
ax2.set_ylabel('Train Loss', fontsize=13)
ax2.set_title('Динаміка втрат під час навчання', fontsize=15, fontweight='bold')
ax2.legend(fontsize=12)
ax2.grid(alpha=0.3)

plt.tight_layout()
plt.show()

# Підсумкова таблиця
print("\n📊 Порівняльна таблиця результатів:\n")
print(f"{'Модель':<25} | {'LR':<8} | {'Фінал Train Acc':<18} | {'Фінал Test Acc':<18} | {'Приріст':<10}")
print("-" * 95)
print(f"{'Baseline (без BN)':<25} | {0.01:<8.2f} | {history_baseline['train_acc'][-1]:<18.2f} | "
      f"{history_baseline['test_acc'][-1]:<18.2f} | {'—':<10}")
print(f"{'CNN + BatchNorm':<25} | {0.1:<8.2f} | {history_bn['train_acc'][-1]:<18.2f} | "
      f"{history_bn['test_acc'][-1]:<18.2f} | {'+' + str(round(history_bn['test_acc'][-1] - history_baseline['test_acc'][-1], 2)) + '%':<10}")

print(f"\n🏆 BatchNorm дозволила:")
print(f"   ✅ Збільшити LR у 10 разів (0.01 → 0.1)")
print(f"   ✅ Покращити test accuracy на {history_bn['test_acc'][-1] - history_baseline['test_acc'][-1]:.2f}%")
print(f"   ✅ Досягти конвергенції швидше (крутіша крива навчання)")