Delete src/bigru_t/training/elastic_weight_consolidation.py

src/bigru_t/training/elastic_weight_consolidation.py

@@ -1,281 +0,0 @@
-"""elastic_weight_consolidation.py — EWC V6.
-
-Adaptado de gru-ring-v13-9-2/source/xavante_work/flexnet/ewc.py.
-
-EWC (Kirkpatrick et al. 2017): previne catastrophic forgetting em treino
-continuado. Mantém Fisher Information Matrix diagonal de uma tarefa anterior
-e penaliza mudanças nos pesos importantes.
-
-Loss total: L = L_task + (λ/2) · Σ_i F_i · (θ_i - θ*_i)²
-
-onde:
-  - F_i = Fisher diagonal = E[(∂log p(y|x) / ∂θ_i)²]
-  - θ* = pesos ótimos da tarefa anterior
-  - λ = peso do regularizador
-
-V6 melhorias vs V13.9.2:
-  * Fisher com clipping logarítmico (preserva ordem, comprime 10^12 → ~10)
-  * Smoothing (moving average) para estabilidade
-  * λ adaptativo (baseado na loss atual)
-  * Correções de 4ª e 6ª ordem Taylor (opcional, default off)
-  * Memory-efficient: Fisher computado em batches (não todo o dataset)
-
-REFERÊNCIAS:
-  - Kirkpatrick et al. (2017). "Overcoming catastrophic forgetting in neural
-    networks". PNAS 114(13):3521-3526.
-  - Huszár (2018). "Note on the quadratic penalties in EWC".
-"""
-from __future__ import annotations
-
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Dict
-from collections import OrderedDict
-
-
-# ═══════════════════════════════════════════════════════════════════════════
-# Fisher Information Matrix diagonal (com clipping logarítmico)
-# ═══════════════════════════════════════════════════════════════════════════
-def compute_fisher_diagonal(
-    model: nn.Module,
-    inputs: torch.Tensor,
-    targets: torch.Tensor,
-    criterion: nn.Module = None,
-    n_samples: int = 1,
-    log_clip: bool = True,
-    clip_scale: float = 1e6,
-) -> Dict[str, torch.Tensor]:
-    """Computa a diagonal da Fisher Information Matrix.
-
-    F_i = E[(∂log p(y|x) / ∂θ_i)²]  ≈  (1/N) · Σ_n (∂L(x_n, y_n) / ∂θ_i)²
-
-    Args:
-        model: modelo
-        inputs: (N, T, d) ou (N, ...) — batch de entradas
-        targets: (N, ...) — targets
-        criterion: função de perda (default: CrossEntropyLoss)
-        n_samples: nº de amostras de Monte Carlo (default 1)
-        log_clip: se True, aplica clipping logarítmico (preserva ordem)
-        clip_scale: escala do clipping (default 1e6)
-
-    Returns:
-        Dict[name → fisher_tensor] com mesmo shape de cada parâmetro
-    """
-    if criterion is None:
-        criterion = nn.CrossEntropyLoss()
-
-    model.eval()  # desativa dropout para Fisher estável
-    fisher: Dict[str, torch.Tensor] = {
-        n: torch.zeros_like(p) for n, p in model.named_parameters()
-        if p.requires_grad
-    }
-
-    for _ in range(n_samples):
-        model.zero_grad()
-        # Forward
-        out = model(inputs)
-        if isinstance(out, tuple):
-            out = out[0]  # modelo pode retornar (logits, aux)
-        # Reshape se necessário
-        if out.dim() > 2:
-            out = out.reshape(-1, out.shape[-1])
-            tgt = targets.reshape(-1)
-        else:
-            tgt = targets
-        loss = criterion(out, tgt)
-        # Backward
-        loss.backward()
-
-        # Acumula (grad)² na Fisher
-        for name, p in model.named_parameters():
-            if p.requires_grad and p.grad is not None:
-                fisher[name] += (p.grad.detach() ** 2) / n_samples
-
-    model.zero_grad()
-
-    # Clipping logarítmico (opcional)
-    if log_clip:
-        for name in fisher:
-            f = fisher[name]
-            # log_clip: F_clipped = log(1 + F / clip_scale) · clip_scale
-            # Preserva ordem (monótono), comprime grandes magnitudes
-            fisher[name] = torch.log1p(f / clip_scale) * clip_scale
-
-    return fisher
-
-
-# ═══════════════════════════════════════════════════════════════════════════
-# ElasticWeightConsolidation — regularizador EWC
-# ═══════════════════════════════════════════════════════════════════════════
-class ElasticWeightConsolidation(nn.Module):
-    """EWC — penaliza mudanças em pesos importantes de tarefas anteriores.
-
-    Args:
-        model: modelo a ser regularizado
-        lambda_init: λ inicial (peso do regularizador) — default 100.0
-        adaptive_lambda: se True, λ é adaptativo (default True)
-        fisher_smoothing: decay do EMA para Fisher (default 0.99)
-        log_clip: se True, aplica clipping logarítmico no Fisher (default True)
-
-    Uso:
-        ewc = ElasticWeightConsolidation(model)
-
-        # Após treinar tarefa A, computa Fisher e θ*
-        ewc.consolidate(inputs_A, targets_A)
-
-        # Durante treino da tarefa B, adiciona loss EWC à loss total
-        loss_total = loss_task + ewc.penalty()
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        lambda_init: float = 100.0,
-        adaptive_lambda: bool = True,
-        fisher_smoothing: float = 0.99,
-        log_clip: bool = True,
-        clip_scale: float = 1e6,
-    ):
-        super().__init__()
-        self.model = model
-        self.fisher_smoothing = fisher_smoothing
-        self.log_clip = log_clip
-        self.clip_scale = clip_scale
-
-        # λ (peso do regularizador) — log-space para positividade
-        self.log_lambda = nn.Parameter(torch.tensor(math.log(lambda_init)))
-        self.adaptive_lambda = adaptive_lambda
-
-        # θ* (pesos ótimos da tarefa anterior) e F (Fisher diagonal)
-        # Inicializados como zeros; serão preenchidos por consolidate()
-        self.star_params: Dict[str, torch.Tensor] = {}
-        self.fisher: Dict[str, torch.Tensor] = {}
-        self.consolidated = False
-
-        # Track da loss anterior para λ adaptativo
-        self.register_buffer("prev_loss", torch.tensor(0.0))
-
-    @property
-    def lambda_val(self) -> torch.Tensor:
-        """λ atual = exp(log_lambda)."""
-        return torch.exp(self.log_lambda)
-
-    @torch.no_grad()
-    def consolidate(
-        self,
-        inputs: torch.Tensor,
-        targets: torch.Tensor,
-        criterion: nn.Module = None,
-        n_samples: int = 1,
-    ) -> dict:
-        """Computa Fisher e θ* após treinar uma tarefa.
-
-        Args:
-            inputs, targets: batch representativo da tarefa concluída
-            criterion: função de perda
-            n_samples: nº de amostras MC para Fisher
-
-        Returns:
-            dict com estatísticas (fisher_mean, fisher_max, n_params)
-        """
-        # 1. Salva θ* (clone dos pesos atuais)
-        self.star_params = {
-            n: p.detach().clone() for n, p in self.model.named_parameters()
-            if p.requires_grad
-        }
-
-        # 2. Computa Fisher diagonal
-        new_fisher = compute_fisher_diagonal(
-            self.model, inputs, targets,
-            criterion=criterion, n_samples=n_samples,
-            log_clip=self.log_clip, clip_scale=self.clip_scale,
-        )
-
-        # 3. Suavização EMA (se já tinha Fisher anterior)
-        if self.consolidated and self.fisher:
-            for name in new_fisher:
-                if name in self.fisher:
-                    self.fisher[name] = (
-                        self.fisher_smoothing * self.fisher[name]
-                        + (1.0 - self.fisher_smoothing) * new_fisher[name]
-                    )
-                else:
-                    self.fisher[name] = new_fisher[name]
-        else:
-            self.fisher = new_fisher
-            self.consolidated = True
-
-        # Estatísticas
-        all_f = torch.cat([f.flatten() for f in self.fisher.values()])
-        return {
-            "fisher_mean": float(all_f.mean().item()),
-            "fisher_max": float(all_f.max().item()),
-            "fisher_min": float(all_f.min().item()),
-            "n_params": int(all_f.numel()),
-            "lambda": float(self.lambda_val.item()),
-        }
-
-    def penalty(self) -> torch.Tensor:
-        """Computa penalidade EWC: (λ/2) · Σ_i F_i · (θ_i - θ*_i)².
-
-        Returns:
-            escalar — penalidade EWC
-        """
-        if not self.consolidated:
-            return torch.tensor(0.0, device=next(self.model.parameters()).device)
-
-        loss = torch.tensor(0.0, device=next(self.model.parameters()).device)
-        lam = self.lambda_val
-
-        for name, p in self.model.named_parameters():
-            if not p.requires_grad:
-                continue
-            if name not in self.fisher or name not in self.star_params:
-                continue
-            F = self.fisher[name].to(p.device)
-            theta_star = self.star_params[name].to(p.device)
-            loss = loss + (F * (p - theta_star) ** 2).sum()
-
-        return 0.5 * lam * loss
-
-    @torch.no_grad()
-    def update_lambda_adaptive(self, current_loss: float) -> float:
-        """Atualiza λ adaptativamente baseado na loss atual.
-
-        Se loss subiu muito (catastrophic forgetting), aumenta λ.
-        Se loss está baixa, mantém λ (não interfere).
-
-        Args:
-            current_loss: loss da tarefa atual
-
-        Returns:
-            novo λ
-        """
-        if not self.adaptive_lambda:
-            return float(self.lambda_val.item())
-
-        prev = float(self.prev_loss.item())
-        # Se loss subiu mais que 50%, aumenta λ (catastrophic forgetting detectado)
-        if current_loss > 1.5 * prev and prev > 0:
-            self.log_lambda.data += 0.1  # λ *= exp(0.1) ≈ 1.10
-        elif current_loss < prev:
-            # Loss caiu — pode reduzir λ gradativamente
-            self.log_lambda.data -= 0.01  # λ /= exp(0.01) ≈ 0.99
-
-        self.prev_loss.fill_(current_loss)
-        return float(self.lambda_val.item())
-
-    def extra_repr(self) -> str:
-        return (
-            f"lambda={float(self.lambda_val.item()):.3f}, "
-            f"consolidated={self.consolidated}, "
-            f"adaptive_lambda={self.adaptive_lambda}"
-        )
-
-
-__all__ = [
-    "compute_fisher_diagonal",
-    "ElasticWeightConsolidation",
-]