Стресс-тест функций активации: GELU vs Tanh

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import time

# --- КОНФИГУРАЦИЯ ---
CONFIG = {
    'INPUT_SIZE': 784,
    'HIDDEN_SIZE': 10,
    'OUTPUT_SIZE': 10,
    'BATCH_SIZE': 512,
    'EPOCHS': 12,
    'LR': 0.003,
    'NUM_RUNS': 20,
    'DEVICE': "cuda" if torch.cuda.is_available() else "cpu",
    
    # Добавили 0.0 (Baseline) во все тесты
    'LOBOTOMY_LEVELS': [0.0, 0.30, 0.50, 0.70, 0.90],
    'SPIKE_LEVELS':    [0.0, 0.30, 0.50, 0.70, 0.90], 
    'NOISE_LEVELS':    [0.0, 0.5, 1.0, 2.0, 3.0] 
}

# --- ЗАГРУЗКА ДАННЫХ ---
class FastMNIST:
    def __init__(self, train=True, device='cpu'):
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
        dataset = datasets.MNIST('./data', train=train, download=True, transform=transform)
        loader = torch.utils.data.DataLoader(dataset, batch_size=len(dataset))
        self.data, self.targets = next(iter(loader))
        self.data = self.data.to(device)
        self.targets = self.targets.to(device)
        self.n_samples = len(self.data)
        
    def get_batches(self, batch_size, shuffle=True):
        if shuffle:
            indices = torch.randperm(self.n_samples, device=self.data.device)
        else:
            indices = torch.arange(self.n_samples, device=self.data.device)
        for start_idx in range(0, self.n_samples, batch_size):
            idx = indices[start_idx : start_idx + batch_size]
            yield self.data[idx], self.targets[idx]

print(f"🔥 DEVICE: {CONFIG['DEVICE']}")
train_data = FastMNIST(train=True, device=CONFIG['DEVICE'])
test_data = FastMNIST(train=False, device=CONFIG['DEVICE'])

# --- МОДЕЛЬ ---
class PrismNet(nn.Module):
    def __init__(self, act_fn, name):
        super().__init__()
        self.name = name
        self.fc1 = nn.Linear(CONFIG['INPUT_SIZE'], CONFIG['HIDDEN_SIZE'])
        self.ln = nn.LayerNorm(CONFIG['HIDDEN_SIZE']) 
        self.act = act_fn()
        self.fc2 = nn.Linear(CONFIG['HIDDEN_SIZE'], CONFIG['OUTPUT_SIZE'])
    
    def forward(self, x, mask=None):
        x = x.view(-1, CONFIG['INPUT_SIZE'])
        pre_latent = self.fc1(x)
        if mask is not None:
            pre_latent = pre_latent * mask
        latent = self.act(self.ln(pre_latent))
        return self.fc2(latent)

# --- ГЕНЕРАТОРЫ МАСОК ---
def get_lobotomy_mask(hidden_size, severity, seed, device):
    torch.manual_seed(seed)
    mask = torch.ones(hidden_size, device=device)
    n_killed = int(hidden_size * severity)
    if n_killed >= hidden_size: n_killed = hidden_size - 1
    
    perm = torch.randperm(hidden_size, device=device)
    killed_indices = perm[:n_killed]
    mask[killed_indices] = 0.0
    
    n_alive = hidden_size - n_killed
    scale = np.sqrt(hidden_size / n_alive)
    return mask * scale

def get_spike_mask(hidden_size, severity, seed, device):
    torch.manual_seed(seed)
    mask = torch.ones(hidden_size, device=device)
    victim = torch.randint(0, hidden_size, (1,)).item()
    
    E_total = float(hidden_size)
    E_spike = severity * E_total
    E_noise = (1.0 - severity) * E_total
    
    amp_spike = np.sqrt(E_spike)
    amp_noise = np.sqrt(E_noise / (hidden_size - 1))
    
    mask[:] = amp_noise
    mask[victim] = amp_spike
    return mask

def get_noise_mask(hidden_size, intensity, seed, device):
    torch.manual_seed(seed)
    raw_noise = torch.randn(hidden_size, device=device) * intensity
    mask = torch.exp(raw_noise) 
    
    current_E = (mask**2).sum()
    target_E = float(hidden_size)
    scale = torch.sqrt(target_E / current_E)
    return mask * scale

# --- ЯДРО ЭКСПЕРИМЕНТА ---
print(f"n=== PRISM FINAL: BASELINE & TRINITY (N={CONFIG['NUM_RUNS']}) ===n")

models_config = [(nn.ReLU, "ReLU"), (nn.Tanh, "Tanh"), (nn.GELU, "GELU")]

# Хранилище
res_lobo = {name: {lvl: [] for lvl in CONFIG['LOBOTOMY_LEVELS']} for _, name in models_config}
res_spike = {name: {lvl: [] for lvl in CONFIG['SPIKE_LEVELS']} for _, name in models_config}
res_noise = {name: {lvl: [] for lvl in CONFIG['NOISE_LEVELS']} for _, name in models_config}

total_start = time.time()

for run in range(CONFIG['NUM_RUNS']):
    run_start = time.time()
    print(f"Run {run+1:02d}/{CONFIG['NUM_RUNS']}...", end=" ", flush=True)
    
    # 1. Train
    trained_models = []
    for act_fn, name in models_config:
        model = PrismNet(act_fn, name).to(CONFIG['DEVICE'])
        opt = optim.Adam(model.parameters(), lr=CONFIG['LR'])
        for epoch in range(CONFIG['EPOCHS']):
            model.train()
            for data, target in train_data.get_batches(CONFIG['BATCH_SIZE']):
                opt.zero_grad()
                logits = model(data)
                loss = nn.CrossEntropyLoss()(logits, target)
                loss.backward()
                opt.step()
        trained_models.append(model)
    
    # Helper for running tests
    def run_test_batch(level_list, result_dict, mask_gen_func):
        for lvl in level_list:
            # Special case for Baseline
            if lvl == 0.0:
                mask = None
            else:
                mask = mask_gen_func(CONFIG['HIDDEN_SIZE'], lvl, seed=1000+run+int(lvl*100), device=CONFIG['DEVICE'])
                
            for model in trained_models:
                model.eval()
                correct = 0; total = 0
                with torch.no_grad():
                    for data, target in test_data.get_batches(2000, shuffle=False):
                        logits = model(data, mask)
                        correct += logits.argmax(1).eq(target).sum().item()
                        total += target.size(0)
                result_dict[model.name][lvl].append(100. * correct / total)

    # 2. Run Tests
    run_test_batch(CONFIG['LOBOTOMY_LEVELS'], res_lobo, get_lobotomy_mask)
    run_test_batch(CONFIG['SPIKE_LEVELS'], res_spike, get_spike_mask)
    run_test_batch(CONFIG['NOISE_LEVELS'], res_noise, get_noise_mask)
            
    print(f"Done ({time.time() - run_start:.1f}s)")

# --- ОТЧЕТ ---
def print_table(title, levels, results_dict, metric_name):
    print(f"nn### {title}")
    print(f"{metric_name:<10} | {'Model':<6} | {'Accuracy':<9} | {'StdDev':<8} | {'95% CI':<16}")
    print("|" + "-"*65 + "|")
    
    for lvl in levels:
        label = str(lvl)
        if lvl == 0.0: label = "0.0 (Base)"
        
        print(f"| **{label}** |        |           |          |                  |")
        for _, name in models_config:
            data = results_dict[name][lvl]
            mean = np.mean(data)
            std = np.std(data)
            ci = 1.96 * std / np.sqrt(len(data))
            
            # Simple highlight logic
            mean_str = f"{mean:.2f}%"
            if lvl > 0.0 and name == "Tanh" and mean > 60: 
                 mean_str = f"**{mean_str}**"
            
            print(f"|            | {name:<6} | {mean_str:<9} | {std:.2f}     | [{mean-ci:.2f}, {mean+ci:.2f}] |")
        print("|" + "-"*65 + "|")

print_table("TEST 1: LOBOTOMY (Потеря информации)", CONFIG['LOBOTOMY_LEVELS'], res_lobo, "Dead %")
print_table("TEST 2: SPIKE (Паразитная доминанта)", CONFIG['SPIKE_LEVELS'], res_spike, "Energy %")
print_table("TEST 3: NOISE (Энтропия / Хаос)", CONFIG['NOISE_LEVELS'], res_noise, "Noise Lvl")

print(f"nTotal Experiment Time: {time.time() - total_start:.1f}s")