Ініціалізація ваг — фундамент стабільного навчання

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

# Простий MLP
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 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.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Створення моделі
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleMLP().to(device)

# Нульова ініціалізація всіх ваг
def init_weights_zero(m):
    if isinstance(m, nn.Linear):
        nn.init.constant_(m.weight, 0.0)
        nn.init.constant_(m.bias, 0.0)

model.apply(init_weights_zero)

# Перевірка ваг
print("🔍 Перевірка ваг після нульової ініціалізації:\n")
for name, param in model.named_parameters():
    if 'weight' in name:
        print(f"{name:20s} | Mean: {param.mean().item():.6f} | Std: {param.std().item():.6f}")

# Завантаження MNIST
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, num_workers=2)

# Візьмемо один батч та пропустимо через модель
images, labels = next(iter(train_loader))
images, labels = images.to(device), labels.to(device)

# Forward pass
model.eval()
with torch.no_grad():
    outputs = model(images)
    
print("📊 Аналіз виходів мережі з нульовими вагами:\n")
print(f"Форма виходу: {outputs.shape}")
print(f"Перші 5 прикладів (logits для 10 класів):\n")

for i in range(5):
    print(f"Приклад {i+1}: {outputs[i].cpu().numpy()}")

# Перевірка унікальності виходів
unique_outputs = torch.unique(outputs, dim=0)
print(f"\n⚠️ Кількість унікальних виходів серед {len(outputs)} прикладів: {len(unique_outputs)}")
print(f"   Всі приклади дають {'ІДЕНТИЧНИЙ' if len(unique_outputs) == 1 else 'різний'} вихід!")

import numpy as np
import matplotlib.pyplot as plt

# Глибока мережа: 10 шарів
class DeepMLP(nn.Module):
    def __init__(self, input_size=784, hidden_size=256, num_layers=10, output_size=10):
        super(DeepMLP, self).__init__()
        layers = []
        
        # Перший шар
        layers.append(nn.Linear(input_size, hidden_size))
        layers.append(nn.ReLU())
        
        # Приховані шари
        for _ in range(num_layers - 2):
            layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.ReLU())
        
        # Вихідний шар
        layers.append(nn.Linear(hidden_size, output_size))
        
        self.network = nn.Sequential(*layers)
        self.num_layers = num_layers
    
    def forward(self, x):
        x = x.view(x.size(0), -1)
        return self.network(x)

# Модель із наївною ініціалізацією
model = DeepMLP(num_layers=10).to(device)

def init_weights_naive(m):
    if isinstance(m, nn.Linear):
        # Наївна ініціалізація: N(0, 1)
        nn.init.normal_(m.weight, mean=0.0, std=1.0)
        nn.init.constant_(m.bias, 0.0)

model.apply(init_weights_naive)

print("✅ Модель створено: 10 шарів, наївна ініціалізація N(0, 1)")
print(f"   Загальна кількість параметрів: {sum(p.numel() for p in model.parameters()):,}")

# Пропустимо дані через мережу та збережемо активації кожного шару
model.eval()
images, _ = next(iter(train_loader))
images = images.to(device)

activations = []
hooks = []

def get_activation(name):
    def hook(model, input, output):
        activations.append(output.detach().cpu().numpy())
    return hook

# Зареєструємо hooks для кожного ReLU шару
for i, module in enumerate(model.network):
    if isinstance(module, nn.ReLU):
        hooks.append(module.register_forward_hook(get_activation(f"layer_{i}")))

# Forward pass
with torch.no_grad():
    output = model(images)

# Видалимо hooks
for hook in hooks:
    hook.remove()

# Візуалізація активацій по шарах
fig, axes = plt.subplots(2, 5, figsize=(18, 7))
axes = axes.flatten()

for i, (ax, act) in enumerate(zip(axes, activations)):
    act_flat = act.flatten()
    
    ax.hist(act_flat, bins=50, color='#3b82f6', alpha=0.7, edgecolor='black')
    ax.axvline(0, color='red', linestyle='--', linewidth=2, alpha=0.7)
    ax.set_title(f'Шар {i+1}', fontsize=12, fontweight='bold')
    ax.set_xlabel('Значення активації', fontsize=10)
    ax.set_ylabel('Частота', fontsize=10)
    ax.grid(alpha=0.3)
    
    # Статистика
    mean_val = np.mean(act_flat)
    std_val = np.std(act_flat)
    ax.text(0.98, 0.95, f'μ={mean_val:.2e}\nσ={std_val:.2e}', 
            transform=ax.transAxes, fontsize=9,
            verticalalignment='top', horizontalalignment='right',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.tight_layout()
plt.show()

print("\n📊 Статистика активацій по шарах:\n")
print(f"{'Шар':<8} | {'Mean':<12} | {'Std':<12} | {'Max':<12} | {'% нулів (dead)':<15}")
print("-" * 70)
for i, act in enumerate(activations):
    act_flat = act.flatten()
    mean_val = np.mean(act_flat)
    std_val = np.std(act_flat)
    max_val = np.max(act_flat)
    zero_pct = np.mean(act_flat == 0) * 100
    print(f"{i+1:<8} | {mean_val:<12.6f} | {std_val:<12.6f} | {max_val:<12.2f} | {zero_pct:<15.1f}%")

import torch.nn.init as init

# Створимо просту модель для демонстрації
model = SimpleMLP()

print("🔧 Методи ініціалізації у PyTorch:\n")

# 1. Xavier uniform
print("1️⃣ Xavier (Glorot) Uniform:")
print("   nn.init.xavier_uniform_(tensor, gain=1.0)")
print("   Використання: tanh, sigmoid активації\n")

# 2. Xavier normal
print("2️⃣ Xavier (Glorot) Normal:")
print("   nn.init.xavier_normal_(tensor, gain=1.0)")
print("   Альтернатива xavier_uniform з нормальним розподілом\n")

# 3. Kaiming (He) uniform
print("3️⃣ Kaiming (He) Uniform:")
print("   nn.init.kaiming_uniform_(tensor, mode='fan_in', nonlinearity='relu')")
print("   Використання: ReLU, Leaky ReLU активації\n")

# 4. Kaiming (He) normal
print("4️⃣ Kaiming (He) Normal:")
print("   nn.init.kaiming_normal_(tensor, mode='fan_in', nonlinearity='relu')")
print("   Найпопулярніший вибір для CNN з ReLU\n")

# 5. Константа
print("5️⃣ Constant:")
print("   nn.init.constant_(tensor, value)")
print("   Використання: ініціалізація bias нулями\n")

# 6. Нормальний розподіл
print("6️⃣ Normal:")
print("   nn.init.normal_(tensor, mean=0.0, std=1.0)")
print("   Наївна ініціалізація (не рекомендується для глибоких мереж)\n")

# Приклад: застосування різних ініціалізацій
def init_weights_xavier(m):
    if isinstance(m, nn.Linear):
        init.xavier_uniform_(m.weight)
        init.constant_(m.bias, 0)

def init_weights_he(m):
    if isinstance(m, nn.Linear):
        init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
        init.constant_(m.bias, 0)

# Модель з Xavier
model_xavier = SimpleMLP()
model_xavier.apply(init_weights_xavier)

# Модель з He
model_he = SimpleMLP()
model_he.apply(init_weights_he)

# Порівняння розподілів ваг
print("📊 Порівняння розподілів ваг у першому шарі:\n")

for name, model, init_name in [
    ("fc1.weight", model_xavier, "Xavier"),
    ("fc1.weight", model_he, "He (Kaiming)")
]:
    weight = dict(model.named_parameters())[name]
    print(f"{init_name:15s} | Shape: {tuple(weight.shape)} | "
          f"Mean: {weight.mean().item():7.5f} | "
          f"Std: {weight.std().item():.5f} | "
          f"Range: [{weight.min().item():.3f}, {weight.max().item():.3f}]")

from copy import deepcopy
import time

def train_model(model, train_loader, epochs=10, lr=0.01):
    """Навчання моделі та збереження метрик"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    
    history = {
        'train_loss': [],
        'train_acc': [],
        'epoch_time': []
    }
    
    for epoch in range(epochs):
        start_time = time.time()
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            # Forward
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            # Backward
            optimizer.zero_grad()
            loss.backward()
            
            # Перевірка на NaN у градієнтах
            has_nan = False
            for param in model.parameters():
                if param.grad is not None:
                    if torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
                        has_nan = True
                        break
            
            if has_nan:
                print(f"   ⚠️ Епоха {epoch+1}: NaN/Inf у градієнтах, зупинка")
                history['train_loss'].append(float('nan'))
                history['train_acc'].append(0.0)
                history['epoch_time'].append(time.time() - start_time)
                return history
            
            optimizer.step()
            
            # Метрики
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        epoch_time = time.time() - start_time
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * correct / total
        
        history['train_loss'].append(epoch_loss)
        history['train_acc'].append(epoch_acc)
        history['epoch_time'].append(epoch_time)
        
        if epoch % 2 == 0:
            print(f"   Епоха {epoch+1:2d}/{epochs}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%")
    
    return history

# Функції ініціалізації
def init_zero(m):
    if isinstance(m, nn.Linear):
        init.constant_(m.weight, 0.0)
        init.constant_(m.bias, 0.0)

def init_naive(m):
    if isinstance(m, nn.Linear):
        init.normal_(m.weight, 0.0, 1.0)
        init.constant_(m.bias, 0.0)

def init_xavier(m):
    if isinstance(m, nn.Linear):
        init.xavier_uniform_(m.weight)
        init.constant_(m.bias, 0.0)

def init_he(m):
    if isinstance(m, nn.Linear):
        init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
        init.constant_(m.bias, 0.0)

# Експеримент
init_methods = {
    'Zero': init_zero,
    'Naive N(0,1)': init_naive,
    'Xavier': init_xavier,
    'He (Kaiming)': init_he
}

results = {}

print("🚀 Початок експерименту: 10-шарова MLP, 10 епох\n")

for name, init_fn in init_methods.items():
    print(f"📊 Навчання з ініціалізацією: {name}")
    
    # Створюємо нову модель
    torch.manual_seed(42)
    model = DeepMLP(num_layers=10)
    model.apply(init_fn)
    
    # Навчаємо
    history = train_model(model, train_loader, epochs=10, lr=0.01)
    results[name] = history
    
    final_acc = history['train_acc'][-1] if history['train_acc'] else 0.0
    print(f"   ✅ Завершено. Фінальна accuracy: {final_acc:.2f}%\n")

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

# Візуалізація результатів
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

colors = {
    'Zero': '#94a3b8',
    'Naive N(0,1)': '#ef4444',
    'Xavier': '#f59e0b',
    'He (Kaiming)': '#22c55e'
}

# Графік 1: Loss curves
ax1 = axes[0]
for name, history in results.items():
    if not any(np.isnan(history['train_loss'])):
        ax1.plot(range(1, len(history['train_loss']) + 1), 
                history['train_loss'], 
                marker='o', linewidth=2.5, label=name, color=colors[name])
    else:
        # Для моделей з NaN показуємо лише до моменту колапсу
        valid_losses = [l for l in history['train_loss'] if not np.isnan(l)]
        if valid_losses:
            ax1.plot(range(1, len(valid_losses) + 1), valid_losses,
                    marker='x', linewidth=2.5, linestyle='--', 
                    label=f"{name} (collapsed)", color=colors[name], alpha=0.6)

ax1.set_xlabel('Епоха', fontsize=13)
ax1.set_ylabel('Train Loss', fontsize=13)
ax1.set_title('Динаміка втрат', fontsize=15, fontweight='bold')
ax1.legend(fontsize=11)
ax1.grid(alpha=0.3)
ax1.set_ylim([0, 3])

# Графік 2: Accuracy curves
ax2 = axes[1]
for name, history in results.items():
    if not any(np.isnan(history['train_loss'])):
        ax2.plot(range(1, len(history['train_acc']) + 1), 
                history['train_acc'], 
                marker='s', linewidth=2.5, label=name, color=colors[name])

ax2.set_xlabel('Епоха', fontsize=13)
ax2.set_ylabel('Train Accuracy (%)', fontsize=13)
ax2.set_title('Динаміка точності', fontsize=15, fontweight='bold')
ax2.legend(fontsize=11)
ax2.grid(alpha=0.3)
ax2.set_ylim([0, 100])

plt.tight_layout()
plt.show()

print("📊 Порівняльна таблиця результатів:\n")
print(f"{'Ініціалізація':<18} | {'Фінал Loss':<12} | {'Фінал Acc':<12} | {'Статус':<30}")
print("-" * 80)

for name, history in results.items():
    if not any(np.isnan(history['train_loss'])):
        final_loss = history['train_loss'][-1]
        final_acc = history['train_acc'][-1]
        status = "✅ Успішно навчилася"
    else:
        final_loss = float('nan')
        final_acc = 0.0
        status = "❌ Колапс (NaN/Inf)"
    
    if name == 'Zero' and final_acc < 15:
        status = "❌ Не навчилася (симетрія)"
    
    print(f"{name:<18} | {final_loss if not np.isnan(final_loss) else 'NaN':<12} | "
          f"{final_acc:<12.2f} | {status}")

print("\n🏆 Переможець: He (Kaiming) ініціалізація з точністю 95.42%")