Hacia el entrenamiento efectivo de redes neuronales

pip install torch torchvision matplotlib tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR
import matplotlib.pyplot as plt
from tqdm import tqdm
import copy

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

# Transforms básicos (los mejoraremos en el paso 7)
transform_base = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2470, 0.2435, 0.2616)),
])

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

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

print(f"Train: {len(train_dataset)} | Test: {len(test_dataset)}")
print(f"Clases: {train_dataset.classes}")
print(f"Input shape: {train_dataset[0][0].shape}")  # → torch.Size([3, 32, 32])

class SimpleCNN(nn.Module):
    """CNN base sin batch norm, sin dropout — para comparar."""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256), nn.ReLU(),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

# Número de parámetros
model = SimpleCNN().to(device)
total_params = sum(p.numel() for p in model.parameters())
print(f"Parámetros: {total_params:,}")  # → ~340K

# ❌ Mala idea: pesos ~ N(0, 1) sin escalar
bad_model = SimpleCNN().to(device)
for p in bad_model.parameters():
    if p.dim() >= 2:
        nn.init.normal_(p, mean=0.0, std=1.0)

# Pasamos un batch y vemos las activaciones
x = torch.randn(4, 3, 32, 32).to(device)
hooks = []
activations = {}
for name, layer in bad_model.features.named_children():
    def hook_fn(n):
        def hook(module, input, output):
            activations[n] = output.detach()
        return hook
    hooks.append(layer.register_forward_hook(hook_fn(name)))

with torch.no_grad():
    bad_model(x)

for name, act in activations.items():
    print(f"Layer {name:>2s}: mean={act.mean():.4f}, std={act.std():.4f}, "
          f"dead={((act==0).float().mean()*100):.1f}%")

for h in hooks: h.remove()

def init_xavier(module):
    """Xavier/Glorot — ideal para tanh, sigmoid."""
    if isinstance(module, (nn.Linear, nn.Conv2d)):
        nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)

model_xavier = SimpleCNN().to(device)
model_xavier.apply(init_xavier)
print("✓ Xavier init aplicada")

def init_kaiming(module):
    """Kaiming/He — ideal para ReLU y variantes."""
    if isinstance(module, (nn.Linear, nn.Conv2d)):
        nn.init.kaiming_normal_(module.weight, mode='fan_in', nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)

model = SimpleCNN().to(device)
model.apply(init_kaiming)
print("✓ Kaiming init aplicada")

# Repetimos el experimento con Kaiming
activations = {}
hooks = []
for name, layer in model.features.named_children():
    def hook_fn(n):
        def hook(module, input, output):
            activations[n] = output.detach()
        return hook
    hooks.append(layer.register_forward_hook(hook_fn(name)))

with torch.no_grad():
    model(x)

for name, act in activations.items():
    print(f"Layer {name:>2s}: mean={act.mean():.4f}, std={act.std():.4f}, "
          f"dead={((act==0).float().mean()*100):.1f}%")

for h in hooks: h.remove()

# SGD puro: w = w - lr * grad
optimizer_sgd = optim.SGD(model.parameters(), lr=0.01)

# Con momentum: acumula velocidad como una "bola rodando cuesta abajo"
# v = momentum * v_prev + grad
# w = w - lr * v
optimizer_sgdm = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Con Nesterov: mira "adelante" antes de calcular el gradiente
optimizer_nesterov = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

# Adam: adapta el learning rate por parámetro usando momentos de 1er y 2do orden
optimizer_adam = optim.Adam(
    model.parameters(),
    lr=1e-3,           # LR base — mucho más bajo que SGD
    betas=(0.9, 0.999), # β1: momentum del gradiente, β2: momentum del gradiente²
    eps=1e-8,           # estabilidad numérica
)

# Valores razonables:
# lr: 1e-3 (default y funciona bien para empezar)
# betas: (0.9, 0.999) — rara vez se cambian
# eps: 1e-8 — puede subirse a 1e-6 si hay inestabilidad numérica

# AdamW: implementa weight decay correctamente (Loshchilov & Hutter, 2019)
# En Adam estándar, weight_decay interactúa con los momentos adaptativos → mal
# AdamW lo desacopla: primero adapta el lr, luego aplica decay por separado
optimizer = optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-2,   # regularización L2 desacoplada (veremos en paso 6)
)
print(f"Optimizer: {optimizer.__class__.__name__}")
print(f"  lr={optimizer.defaults['lr']}, wd={optimizer.defaults['weight_decay']}")

def get_param_groups(model, lr=1e-3, weight_decay=1e-2):
    """Separa parámetros que deben tener weight decay de los que no."""
    decay_params = []
    no_decay_params = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # No aplicar weight decay a bias ni a parámetros de normalización
        if param.dim() <= 1 or 'bias' in name:
            no_decay_params.append(param)
        else:
            decay_params.append(param)
    return [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]

param_groups = get_param_groups(model)
optimizer = optim.AdamW(param_groups, lr=1e-3)

n_decay = sum(p.numel() for p in param_groups[0]["params"])
n_nodecay = sum(p.numel() for p in param_groups[1]["params"])
print(f"Con weight decay: {n_decay:,} params | Sin weight decay: {n_nodecay:,} params")

from torch.optim.lr_scheduler import StepLR

# Reduce el LR por un factor gamma cada step_size épocas
scheduler_step = StepLR(optimizer, step_size=30, gamma=0.1)
# Epoch  0-29: lr = 1e-3
# Epoch 30-59: lr = 1e-4
# Epoch 60-89: lr = 1e-5

# En el training loop:
# for epoch in range(num_epochs):
#     train_one_epoch(...)
#     scheduler_step.step()  ← ¡al final de cada época!

from torch.optim.lr_scheduler import CosineAnnealingLR

# Decae el LR siguiendo medio coseno: de lr_max a eta_min
scheduler_cosine = CosineAnnealingLR(
    optimizer,
    T_max=100,       # número total de épocas
    eta_min=1e-6,    # LR mínimo al final
)
# Suave y sin escalones bruscos → muy popular en visión por computador

from torch.optim.lr_scheduler import OneCycleLR

# Super-convergence: sube el LR rápidamente y luego baja
# (Smith & Topin, 2018: "Super-Convergence")
scheduler_onecycle = OneCycleLR(
    optimizer,
    max_lr=1e-2,                # pico máximo del LR
    epochs=100,
    steps_per_epoch=len(train_loader),  # ← se actualiza POR BATCH, no por época
    pct_start=0.3,              # 30% del training para subir, 70% para bajar
    anneal_strategy='cos',
    div_factor=10,              # lr_inicial = max_lr / div_factor
    final_div_factor=100,       # lr_final = lr_inicial / final_div_factor
)

# ¡IMPORTANTE! OneCycleLR se llama por batch, NO por época:
# for epoch in range(num_epochs):
#     for batch in train_loader:
#         ...
#         optimizer.step()
#         scheduler_onecycle.step()  ← después de cada batch

import math

class WarmupCosineScheduler:
    """LR warmup lineal + cosine decay. Estilo Transformers."""
    def __init__(self, optimizer, warmup_epochs, total_epochs, min_lr=1e-6):
        self.optimizer = optimizer
        self.warmup = warmup_epochs
        self.total = total_epochs
        self.min_lr = min_lr
        self.base_lrs = [g['lr'] for g in optimizer.param_groups]
        self.step_count = 0

    def step(self):
        self.step_count += 1
        if self.step_count <= self.warmup:
            # Warmup lineal
            scale = self.step_count / self.warmup
        else:
            # Cosine decay
            progress = (self.step_count - self.warmup) / (self.total - self.warmup)
            scale = 0.5 * (1 + math.cos(math.PI * progress))
        for g, base_lr in zip(self.optimizer.param_groups, self.base_lrs):
            g['lr'] = self.min_lr + (base_lr - self.min_lr) * scale

    def get_lr(self):
        return [g['lr'] for g in self.optimizer.param_groups]

# Uso
scheduler = WarmupCosineScheduler(optimizer, warmup_epochs=5, total_epochs=100)

# Después de loss.backward(), antes de optimizer.step():
max_norm = 1.0  # umbral de la norma L2

# Calcula la norma total del gradiente y reescala si supera max_norm
total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)
print(f"Grad norm: {total_norm:.4f} (clipped to {max_norm})")

# Alternativa: clip por valor (menos común)
# nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

class CNN_BN(nn.Module):
    """CNN con Batch Normalization entre Conv y ReLU."""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),         # ← BN para mapas 2D
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.BatchNorm1d(256),        # ← BN para vectores 1D
            nn.ReLU(),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

# LayerNorm: normaliza sobre features, no sobre el batch
# → Estándar en Transformers (independiente del batch size)
ln = nn.LayerNorm(normalized_shape=256)

# GroupNorm: divide canales en G grupos, normaliza cada grupo
# → Útil cuando batch size es muy pequeño (batch=1–4)
gn = nn.GroupNorm(num_groups=8, num_channels=64)

# InstanceNorm: normaliza cada mapa de features individualmente
# → Estándar en style transfer
inst = nn.InstanceNorm2d(64)

def get_grad_norm(model):
    """Calcula la norma L2 total del gradiente."""
    total_norm = 0.0
    for p in model.parameters():
        if p.grad is not None:
            total_norm += p.grad.data.norm(2).item() ** 2
    return total_norm ** 0.5

# Usar en el training loop:
# loss.backward()
# grad_norm = get_grad_norm(model)
# if grad_norm > 100:
#     print(f"⚠️ Grad norm alta: {grad_norm:.1f}")
# nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# optimizer.step()

# Weight decay típico: 1e-2 para grandes modelos, 1e-4 para CNNs
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4   # ← Regularización L2 desacoplada
)

# NOTA: En SGD, weight_decay=1e-4 es equivalente a L2 clásico.
# En AdamW, está desacoplado del cálculo de momentos (mejor).

class CNN_Regularized(nn.Module):
    """CNN con BatchNorm + Dropout."""
    def __init__(self, num_classes=10, dropout_rate=0.3):
        super().__init__()
        self.features = nn.Sequential(
            # Bloque 1
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.1),          # ← Dropout espacial (ligero en features)

            # Bloque 2
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.15),

            # Bloque 3
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(dropout_rate),   # ← Dropout estándar en FC
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

model = CNN_Regularized(num_classes=10, dropout_rate=0.3)
print(f"Parámetros: {sum(p.numel() for p in model.parameters()):,}")

# PyTorch ≥ 1.10 soporta label_smoothing directamente
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# Equivale a:
# y_smooth = (1 - ε) * one_hot(y) + ε / num_classes
# loss = CrossEntropy(logits, y_smooth)

# Ejemplo: para clase 3 con 10 clases y ε=0.1:
# [0.01, 0.01, 0.01, 0.91, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01]

class EarlyStopping:
    """Detiene el entrenamiento si no mejora tras `patience` épocas."""
    def __init__(self, patience=5, min_delta=1e-4):
        self.patience = patience
        self.min_delta = min_delta
        self.best_score = None
        self.counter = 0
        self.should_stop = False

    def step(self, val_metric):
        """Llamar tras cada época con la métrica de validación (mayor = mejor)."""
        if self.best_score is None:
            self.best_score = val_metric
        elif val_metric < self.best_score + self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.should_stop = True
                print(f"⏹ Early stopping: sin mejora en {self.patience} épocas.")
        else:
            self.best_score = val_metric
            self.counter = 0

# Uso:
# early_stop = EarlyStopping(patience=7)
# for epoch in range(max_epochs):
#     train(...)
#     val_acc = evaluate(...)
#     early_stop.step(val_acc)
#     if early_stop.should_stop:
#         break

from torchvision import transforms

# ── Pipeline de entrenamiento CON augmentation ──────────────
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),          # Recorte aleatorio con padding
    transforms.RandomHorizontalFlip(p=0.5),        # Flip horizontal 50%
    transforms.ColorJitter(
        brightness=0.2, contrast=0.2,
        saturation=0.2, hue=0.1
    ),                                             # Variación de color
    transforms.RandomRotation(15),                 # Rotación ±15°
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2470, 0.2435, 0.2616]
    ),
    transforms.RandomErasing(p=0.25, scale=(0.02, 0.2)),  # "Cutout" aleatorio
])

# ── Pipeline de validación/test SIN augmentation ────────────
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.4914, 0.4822, 0.4465],
        std=[0.2470, 0.2435, 0.2616]
    ),
])

# Aplicar a los datasets
train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, transform=transform_train, download=True
)
test_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, transform=transform_test, download=True
)

import numpy as np

def mixup_batch(images, targets, alpha=0.2):
    """
    Mixup (Zhang et al., 2018): mezcla pares de imágenes y sus labels.
    images: (B, C, H, W), targets: (B,) con índices de clase
    Retorna: images mezcladas, targets_a, targets_b, lambda
    """
    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1.0
    batch_size = images.size(0)
    index = torch.randperm(batch_size, device=images.device)

    mixed_images = lam * images + (1 - lam) * images[index]
    return mixed_images, targets, targets[index], lam

def mixup_criterion(criterion, pred, y_a, y_b, lam):
    """Loss ponderada entre las dos clases mezcladas."""
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

# En el training loop:
# images, targets = batch
# images, targets_a, targets_b, lam = mixup_batch(images, targets, alpha=0.2)
# outputs = model(images)
# loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)

def cutmix_batch(images, targets, alpha=1.0):
    """
    CutMix (Yun et al., 2019): pega un parche de otra imagen en cada imagen.
    Más efectivo que Mixup para localización y detección.
    """
    lam = np.random.beta(alpha, alpha)
    batch_size = images.size(0)
    index = torch.randperm(batch_size, device=images.device)

    _, _, H, W = images.shape

    # Calcular el tamaño del parche a recortar
    cut_ratio = np.sqrt(1.0 - lam)
    cut_h = int(H * cut_ratio)
    cut_w = int(W * cut_ratio)

    # Centro aleatorio del parche
    cy = np.random.randint(H)
    cx = np.random.randint(W)

    # Coordenadas del bounding box
    y1 = np.clip(cy - cut_h // 2, 0, H)
    y2 = np.clip(cy + cut_h // 2, 0, H)
    x1 = np.clip(cx - cut_w // 2, 0, W)
    x2 = np.clip(cx + cut_w // 2, 0, W)

    # Pegar el parche de la imagen permutada
    mixed = images.clone()
    mixed[:, :, y1:y2, x1:x2] = images[index, :, y1:y2, x1:x2]

    # Ajustar lambda al área real del parche
    lam = 1 - ((y2 - y1) * (x2 - x1)) / (H * W)
    return mixed, targets, targets[index], lam

# AutoAugment: política aprendida por RL (Cubuk et al., 2019)
transform_auto = transforms.Compose([
    transforms.AutoAugment(transforms.AutoAugmentPolicy.CIFAR10),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
])

# TrivialAugment: más simple, sin búsqueda, ¡y funciona igual o mejor!
# (Müller & Hutter, 2021)
transform_trivial = transforms.Compose([
    transforms.TrivialAugmentWide(),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
])

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

SEED = 42
torch.manual_seed(SEED)

# ── Transformaciones ─────────────────────────────────────────
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
    transforms.RandomErasing(p=0.25),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]),
])

# ── Datasets y DataLoaders ───────────────────────────────────
full_train = torchvision.datasets.CIFAR10('./data', train=True, transform=transform_train, download=True)
train_set, val_set = random_split(full_train, [45000, 5000],
                                   generator=torch.Generator().manual_seed(SEED))
# Validación con transforms limpios
val_set.dataset = torchvision.datasets.CIFAR10('./data', train=True, transform=transform_test)
test_set = torchvision.datasets.CIFAR10('./data', train=False, transform=transform_test)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True,  num_workers=2, pin_memory=True)
val_loader   = DataLoader(val_set,   batch_size=256, shuffle=False, num_workers=2, pin_memory=True)
test_loader  = DataLoader(test_set,  batch_size=256, shuffle=False, num_workers=2, pin_memory=True)

print(f"Train: {len(train_set)} | Val: {len(val_set)} | Test: {len(test_set)}")

class EfficientCNN(nn.Module):
    """
    CNN mejorada para CIFAR-10 con:
    ✅ Kaiming initialization
    ✅ BatchNorm en todas las capas
    ✅ Dropout espacial (features) + estándar (classifier)
    """
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            # Bloque 1: 32×32 → 16×16
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.1),

            # Bloque 2: 16×16 → 8×8
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.15),

            # Bloque 3: 8×8 → 4×4
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes),
        )

        # ✅ Kaiming initialization
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                nn.init.zeros_(m.bias)

    def forward(self, x):
        return self.classifier(self.features(x))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = EfficientCNN(num_classes=10).to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"EfficientCNN: {n_params:,} parámetros | device: {device}")

# ── Parameter groups (no weight decay en bias/BN) ──────────
def get_param_groups(model, base_wd=1e-4):
    decay, no_decay = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.ndim <= 1 or 'bias' in name:  # bias + BN params
            no_decay.append(param)
        else:
            decay.append(param)
    return [
        {'params': decay, 'weight_decay': base_wd},
        {'params': no_decay, 'weight_decay': 0.0},
    ]

NUM_EPOCHS = 50

# ✅ AdamW con weight decay desacoplado
optimizer = torch.optim.AdamW(get_param_groups(model), lr=1e-3)

# ✅ OneCycleLR: warm-up + cosine decay automático
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-3,
    epochs=NUM_EPOCHS,
    steps_per_epoch=len(train_loader),
    pct_start=0.1,      # 10% warm-up
    anneal_strategy='cos',
)

# ✅ Label smoothing
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# ✅ Early stopping
class EarlyStopping:
    def __init__(self, patience=7, min_delta=1e-4):
        self.patience, self.min_delta = patience, min_delta
        self.best_score, self.counter, self.should_stop = None, 0, False
    def step(self, val_metric):
        if self.best_score is None or val_metric > self.best_score + self.min_delta:
            self.best_score, self.counter = val_metric, 0
        else:
            self.counter += 1
            self.should_stop = self.counter >= self.patience

early_stop = EarlyStopping(patience=7)
MAX_GRAD_NORM = 1.0  # ✅ Gradient clipping

def train_one_epoch(model, loader, optimizer, scheduler, criterion, device, max_grad_norm):
    model.train()
    total_loss, correct, total = 0.0, 0, 0

    for images, targets in loader:
        images, targets = images.to(device), targets.to(device)

        # Forward
        outputs = model(images)
        loss = criterion(outputs, targets)

        # Backward
        optimizer.zero_grad()
        loss.backward()

        # ✅ Gradient clipping
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)

        optimizer.step()
        scheduler.step()  # OneCycleLR: step por batch

        # Métricas
        total_loss += loss.item() * targets.size(0)
        correct += (outputs.argmax(dim=1) == targets).sum().item()
        total += targets.size(0)

    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0

    for images, targets in loader:
        images, targets = images.to(device), targets.to(device)
        outputs = model(images)
        loss = criterion(outputs, targets)
        total_loss += loss.item() * targets.size(0)
        correct += (outputs.argmax(dim=1) == targets).sum().item()
        total += targets.size(0)

    return total_loss / total, correct / total

import time

history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': [], 'lr': []}
best_val_acc = 0.0

print(f"{'Epoch':>5} {'Train Loss':>11} {'Train Acc':>10} {'Val Loss':>10} {'Val Acc':>9} {'LR':>10}")
print("─" * 62)

for epoch in range(1, NUM_EPOCHS + 1):
    t0 = time.time()

    train_loss, train_acc = train_one_epoch(
        model, train_loader, optimizer, scheduler, criterion, device, MAX_GRAD_NORM
    )
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)

    current_lr = optimizer.param_groups[0]['lr']

    # Guardar historial
    history['train_loss'].append(train_loss)
    history['train_acc'].append(train_acc)
    history['val_loss'].append(val_loss)
    history['val_acc'].append(val_acc)
    history['lr'].append(current_lr)

    # Guardar mejor modelo
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), 'best_model.pth')

    elapsed = time.time() - t0
    print(f"{epoch:5d} {train_loss:11.4f} {train_acc:9.2%} {val_loss:10.4f} {val_acc:8.2%} {current_lr:10.6f}")

    # ✅ Early stopping
    early_stop.step(val_acc)
    if early_stop.should_stop:
        print(f"\n⏹ Early stopping en época {epoch}")
        break

print(f"\n🏆 Mejor val accuracy: {best_val_acc:.2%}")

# Cargar el mejor modelo
model.load_state_dict(torch.load('best_model.pth', weights_only=True))
test_loss, test_acc = evaluate(model, test_loader, criterion, device)

print(f"📊 Test accuracy: {test_acc:.2%}")
print(f"📊 Test loss:     {test_loss:.4f}")

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# Loss
axes[0].plot(history['train_loss'], label='Train')
axes[0].plot(history['val_loss'], label='Val')
axes[0].set_xlabel('Época'); axes[0].set_ylabel('Loss')
axes[0].set_title('Loss'); axes[0].legend()

# Accuracy
axes[1].plot(history['train_acc'], label='Train')
axes[1].plot(history['val_acc'], label='Val')
axes[1].set_xlabel('Época'); axes[1].set_ylabel('Accuracy')
axes[1].set_title('Accuracy'); axes[1].legend()

# Learning rate
axes[2].plot(history['lr'], color='tab:orange')
axes[2].set_xlabel('Época'); axes[2].set_ylabel('LR')
axes[2].set_title('Learning Rate')

plt.tight_layout()
plt.savefig('training_curves.png', dpi=150)
plt.show()

def log_grad_norms(model):
    """Imprime la norma del gradiente de cada capa con parámetros."""
    rows = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.data.norm(2).item()
            rows.append((name, grad_norm, param.data.norm(2).item()))

    print(f"{'Capa':<40} {'|grad|':>10} {'|peso|':>10} {'ratio':>10}")
    print("─" * 72)
    for name, gn, wn in rows:
        ratio = gn / (wn + 1e-8)
        flag = " ⚠️" if gn > 10 or gn < 1e-7 else ""
        print(f"{name:<40} {gn:>10.4f} {wn:>10.4f} {ratio:>10.6f}{flag}")

# Usar después de loss.backward():
# log_grad_norms(model)