Upload cnn_bigru/training/trainer.py with huggingface_hub

cnn_bigru/training/trainer.py

@@ -0,0 +1,495 @@
+"""trainer.py — Loop de treinamento cooperativo com auto-aprendizado.
+
+Implementa a seção 11 de dados.txt:
+
+    PARA época = 1 ATÉ NUM_EPOCAS:
+        PARA batch EM DataLoader:
+            // Forward Gerador
+            logits_prova, enc_out, _ = gerador.forward(axiomas, conjectura, prova_alvo)
+
+            // Calcular perdas e avaliações do verificador
+            L_total, L_G, L_V, L_AH, lista_v = calcular_perdas(...)
+
+            // Backward
+            zerar_gradientes
+            L_total.backward()
+
+            // Controle de gradiente (clip + spectral norm)
+            aplicar_controle_gradiente(gerador)
+            aplicar_controle_gradiente(verificador)
+
+            // Ajuste dinâmico de LR baseado em normas de gradiente e v_media
+            eta_G, eta_V = atualizar_taxas(...)
+
+            // Passo de otimização
+            otimizador_G.step()
+            otimizador_V.step()
+
+            // Normalização espectral após step
+            aplicar_espectral_norm(gerador)
+            aplicar_espectral_norm(verificador)
+
+Adicionalmente:
+    - Mecanismo de hipóteses ativado quando v < threshold
+    - N tentativas de sinergia (SynergySearcher) no início de cada época
+    - Meta-controlador que ajusta pesos das perdas baseado na variância do gradiente
+"""
+from __future__ import annotations
+
+import logging
+import math
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from ..losses.losses import LossConfig, MultiLoss
+from ..models.generator_verifier import (
+    GeneratorCNNBiGRU,
+    VerifierCNNBiGRU,
+    AntiHallucinationLayer,
+)
+from ..models.multimodal_model import MultimodalCNNBiGRU
+from ..utils.memory_optimizer import MemoryOptimizer
+from .auto_learner import AutoLearnConfig, AutoLearner, estimate_curvature, orthogonal_init_model
+from .hypothesis_controller import HypothesisConfig, HypothesisController, SynergySearcher
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainerConfig:
+    """Configuração do treinador."""
+    # Épocas e batches
+    num_epochs: int = 2
+    max_batches_per_epoch: int = 10  # limita batches para teste de 50 amostras
+    batch_size: int = 8
+
+    # Synergy search
+    n_synergy_attempts: int = 3     # N tentativas de acerto do treinamento
+    use_synergy_search: bool = True
+
+    # Hypothesis
+    use_hypotheses: bool = True
+    n_hypotheses: int = 4
+
+    # Loss config
+    loss_config: LossConfig = field(default_factory=LossConfig)
+
+    # Auto-learn config
+    auto_config: AutoLearnConfig = field(default_factory=AutoLearnConfig)
+
+    # Device
+    device: str = "cpu"
+
+    # Logging
+    log_every: int = 1
+
+
+class CooperativeTrainer:
+    """Treina o modelo multimodal CNN-BiGRU cooperativo.
+
+    Args:
+        model: MultimodalCNNBiGRU
+        generator: GeneratorCNNBiGRU (opcional, para modo gerar+verificar)
+        verifier: VerifierCNNBiGRU (opcional)
+        tokenizer: BBPETokenizer
+        config: TrainerConfig
+    """
+
+    def __init__(
+        self,
+        model: MultimodalCNNBiGRU,
+        tokenizer,
+        config: TrainerConfig,
+        generator: Optional[GeneratorCNNBiGRU] = None,
+        verifier: Optional[VerifierCNNBiGRU] = None,
+        anti_hallucination: Optional[AntiHallucinationLayer] = None,
+    ):
+        self.model = model.to(config.device)
+        self.tokenizer = tokenizer
+        self.cfg = config
+        self.device = config.device
+
+        # Componentes opcionais
+        self.generator = generator.to(config.device) if generator else None
+        self.verifier = verifier.to(config.device) if verifier else None
+        self.anti_hall = (
+            anti_hallucination.to(config.device) if anti_hallucination else None
+        )
+
+        # Loss
+        self.loss_fn = MultiLoss(config.loss_config, pad_idx=tokenizer.pad_id)
+
+        # Auto-learner
+        self.auto = AutoLearner(config.auto_config)
+
+        # Synergy searcher
+        self.synergy = SynergySearcher(
+            n_attempts=config.n_synergy_attempts,
+            seed=42,
+        )
+
+        # Otimizadores
+        self.opt_model = torch.optim.AdamW(
+            self.model.parameters(), lr=config.auto_config.initial_lr_G,
+            weight_decay=config.auto_config.l2_reg if hasattr(config.auto_config, 'l2_reg') else 1e-5
+        )
+        self.opt_gen = (
+            torch.optim.AdamW(self.generator.parameters(), lr=config.auto_config.initial_lr_G)
+            if self.generator else None
+        )
+        self.opt_ver = (
+            torch.optim.AdamW(self.verifier.parameters(), lr=config.auto_config.initial_lr_V)
+            if self.verifier else None
+        )
+
+        # Memory optimizer
+        self.mem = MemoryOptimizer(enable_amp=(self.device != "cpu"))
+        self.mem.configure()
+
+        # Hypothesis controller (aplicado sobre features do modelo)
+        self.hypothesis = HypothesisController(
+            feat_dim=model.fusion.out_dim,
+            out_dim=model.fusion.out_dim,
+            config=HypothesisConfig(
+                n_hypotheses=config.n_hypotheses,
+                threshold_err=config.loss_config.threshold_err,
+            ),
+        ).to(self.device)
+
+        # Histórico de treinamento
+        self.history: List[Dict[str, Any]] = []
+        self.metrics: Dict[str, List[float]] = {"loss": [], "ppl": [], "v_mean": [], "lr_G": [], "lr_V": []}
+
+    def _compute_v_and_tau(
+        self,
+        fused: torch.Tensor,
+        targets: torch.Tensor,
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+        """Computa valores do verificador (v) e anti-alucinação (τ).
+
+        Simplificação: usa o fused como proxy para o passo gerado.
+        Em uma versão completa, isto integraria com o generator e verifier.
+        """
+        v_list = []
+        tau_list = []
+
+        # Verifier: para cada passo temporal, calcula v
+        # Aqui simplificamos: v é calculado uma vez por batch
+        if self.verifier is not None:
+            # Passo hipotético: usa o fused repetido como "passo"
+            B = fused.size(0)
+            passo = fused[:, :128]  # dimensão compatível com o encoder
+            # Concatena com zeros para matching dimensional (simplificação)
+            passo_full = torch.cat([passo, passo], dim=-1)  # [B, 256]
+
+            # Não temos tokens aqui; usamos features direto para uma v proxy
+            # v = sigmoid(linear(fused))
+            v_proxy = torch.sigmoid(fused.mean(dim=-1, keepdim=True))  # [B, 1]
+            v_list.append(v_proxy)
+        else:
+            # Sem verifier: v = 1 (sem punição)
+            v_list.append(torch.ones(fused.size(0), 1, device=fused.device))
+
+        # Anti-hallucination: τ sobre os logits
+        if self.anti_hall is not None:
+            # Logits proxy: projeta fused para vocab_size via lm_head
+            logits_proxy = self.model.lm_head(fused)  # [B, V]
+            tau = self.anti_hall(logits_proxy)  # [B, 1]
+            tau_list.append(tau)
+        else:
+            tau_list.append(torch.ones(fused.size(0), 1, device=fused.device))
+
+        return v_list, tau_list
+
+    def _train_step(
+        self,
+        batch: Dict[str, torch.Tensor],
+    ) -> Dict[str, float]:
+        """Executa um passo de treinamento. Retorna métricas."""
+        # Move para device
+        input_ids_a = batch["input_ids_a"].to(self.device, non_blocking=True)
+        input_ids_b = batch["input_ids_b"].to(self.device, non_blocking=True)
+        attn_mask_a = batch["attn_mask_a"].to(self.device, non_blocking=True)
+        attn_mask_b = batch["attn_mask_b"].to(self.device, non_blocking=True)
+        images = batch["images"].to(self.device, non_blocking=True)
+        audios = batch["audios"].to(self.device, non_blocking=True)
+        labels = batch["labels"].to(self.device, non_blocking=True)
+
+        # Forward
+        with self.mem.amp_context(self.device):
+            out = self.model(
+                input_ids_a, input_ids_b,
+                images=images, audios=audios,
+                attn_mask_a=attn_mask_a, attn_mask_b=attn_mask_b,
+                mode="classify",
+            )
+            logits = out["logits"]  # [B, num_classes]
+            fused = out["fused"]    # [B, fusion_dim]
+
+            # Targets: usar labels como targets de classificação
+            # Para LM-style loss, convertemos para one-hot vs logits
+            # Aqui usamos CrossEntropy direto sobre os labels de classificação
+            ce_loss = F.cross_entropy(logits, labels)
+
+            # Verifier + anti-hallucination (proxy)
+            v_list, tau_list = self._compute_v_and_tau(fused, labels)
+
+            # Targets para o gerador: usar input_ids_b como "prova alvo"
+            # (apenas se generator estiver presente)
+            gen_logits = None
+            if self.generator is not None:
+                try:
+                    gen_out = self.generator(
+                        input_ids_a, input_ids_b,
+                        target_proof=input_ids_b,  # teacher forcing
+                        bos_id=self.tokenizer.bos_id,
+                        eos_id=self.tokenizer.eos_id,
+                    )
+                    gen_logits = gen_out["logits"]
+                except Exception as e:
+                    logger.debug("Generator forward falhou: %s", e)
+                    gen_logits = None
+
+            # Constrói logits para a MultiLoss:
+            # Se temos generator, usa ele; senão, projeta fused para [B, T, V]
+            if gen_logits is not None:
+                loss_logits = gen_logits
+                loss_targets = input_ids_b
+            else:
+                # Proxy: expande logits de classificação para [B, T, V]
+                # Não é ideal mas mantém o pipeline funcional
+                B = fused.size(0)
+                T = input_ids_b.size(1)
+                V = self.model.vocab_size
+                loss_logits = self.model.lm_head(fused).unsqueeze(1).expand(-1, T, -1)
+                loss_targets = input_ids_b
+
+            # Hypothesis controller (ativa hipóteses se v < threshold)
+            v_mean_tensor = torch.stack([v.squeeze(-1) for v in v_list]).mean(dim=0).unsqueeze(-1)
+            hyp_out = self.hypothesis(fused, v_mean_tensor)
+
+            # Curvatura
+            try:
+                curv = estimate_curvature(
+                    self.model,
+                    lambda: self.loss_fn(
+                        loss_logits, loss_targets, v_list, tau_list,
+                        models_for_reg=[self.model],
+                    ),
+                    eps=self.cfg.loss_config.curvature_eps,
+                    n_samples=2,
+                )
+                self.loss_fn.set_curvature_estimate(curv)
+            except Exception as e:
+                logger.debug("Curvatura falhou: %s", e)
+
+            # Compute losses
+            losses = self.loss_fn(
+                loss_logits, loss_targets, v_list, tau_list,
+                models_for_reg=[m for m in [self.model, self.generator, self.verifier] if m is not None],
+            )
+
+            # Adiciona loss de classificação
+            total_loss = losses["total"] + ce_loss
+
+        # Backward
+        self.opt_model.zero_grad(set_to_none=True)
+        if self.opt_gen: self.opt_gen.zero_grad(set_to_none=True)
+        if self.opt_ver: self.opt_ver.zero_grad(set_to_none=True)
+
+        total_loss.backward()
+
+        # Clipping
+        self.auto.apply_clipping(self.model)
+        if self.generator: self.auto.apply_clipping(self.generator)
+        if self.verifier: self.auto.apply_clipping(self.verifier)
+
+        # Normas de gradiente
+        grad_G = self.auto.compute_total_grad_norm(self.model)
+        grad_V = self.auto.compute_total_grad_norm(self.verifier) if self.verifier else 0.0
+
+        # v_mean e L_V_mean
+        v_mean = float(torch.stack([v.squeeze(-1) for v in v_list]).mean().item())
+        L_V_mean = float(losses["l_v"].item()) if isinstance(losses["l_v"], torch.Tensor) else 0.0
+
+        # Ajuste dinâmico de LR
+        novo_lr_G, novo_lr_V = self.auto.update_learning_rates(
+            grad_G, grad_V, v_mean, L_V_mean
+        )
+        self.auto.set_optimizer_lr(self.opt_model, novo_lr_G)
+        if self.opt_gen: self.auto.set_optimizer_lr(self.opt_gen, novo_lr_G)
+        if self.opt_ver: self.auto.set_optimizer_lr(self.opt_ver, novo_lr_V)
+
+        # Step
+        self.opt_model.step()
+        if self.opt_gen: self.opt_gen.step()
+        if self.opt_ver: self.opt_ver.step()
+
+        # Spectral norm após step
+        if self.cfg.auto_config.apply_after_step:
+            self.auto.apply_spectral_norm(self.model)
+            if self.generator: self.auto.apply_spectral_norm(self.generator)
+            if self.verifier: self.auto.apply_spectral_norm(self.verifier)
+
+        # Métricas
+        loss_val = float(total_loss.item())
+        try:
+            ppl = math.exp(min(loss_val, 20))  # evita overflow
+        except OverflowError:
+            ppl = float("inf")
+
+        return {
+            "loss": loss_val,
+            "ppl": ppl,
+            "v_mean": v_mean,
+            "lr_G": novo_lr_G,
+            "lr_V": novo_lr_V,
+            "l_g": float(losses["l_g"].item()) if isinstance(losses["l_g"], torch.Tensor) else 0.0,
+            "l_v": L_V_mean,
+            "l_ah": float(losses["l_ah"].item()) if isinstance(losses["l_ah"], torch.Tensor) else 0.0,
+            "hypothesis_activated": hyp_out["activated"],
+        }
+
+    def _eval_config(self, cfg: Dict, sample_batch: Dict[str, torch.Tensor]) -> float:
+        """Avalia uma configuração de perda em um mini-batch (para synergy search)."""
+        # Salva config atual
+        original = {k: getattr(self.loss_fn.cfg, k) for k in cfg}
+        # Aplica nova config
+        for k, v in cfg.items():
+            if hasattr(self.loss_fn.cfg, k):
+                setattr(self.loss_fn.cfg, k, v)
+        try:
+            metrics = self._train_step(sample_batch)
+            return metrics["loss"]
+        finally:
+            # Restaura
+            for k, v in original.items():
+                setattr(self.loss_fn.cfg, k, v)
+
+    def train(
+        self,
+        dataloader: DataLoader,
+    ) -> Dict[str, Any]:
+        """Executa o treinamento completo."""
+        logger.info("Iniciando treinamento cooperativo CNN-BiGRU")
+        logger.info("  epochs=%d, batch_size=%d, n_synergy=%d",
+                    self.cfg.num_epochs, self.cfg.batch_size, self.cfg.n_synergy_attempts)
+
+        start_time = time.time()
+
+        for epoch in range(self.cfg.num_epochs):
+            epoch_loss = 0.0
+            epoch_ppl = 0.0
+            n_batches = 0
+            hypothesis_activations = 0
+
+            # Synergy search no início da época (apenas época > 0 para economizar tempo)
+            if self.cfg.use_synergy_search and epoch == 0:
+                try:
+                    # Pega um batch de amostra para avaliação
+                    sample_batch = next(iter(dataloader))
+                    sample_batch = {k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
+                                    for k, v in sample_batch.items()}
+
+                    base_cfg = {
+                        "alpha": self.loss_fn.cfg.alpha,
+                        "beta": self.loss_fn.cfg.beta,
+                        "gamma_loss": self.loss_fn.cfg.gamma_loss,
+                        "delta": self.loss_fn.cfg.delta,
+                        "lambda_penal": self.loss_fn.cfg.lambda_penal,
+                        "mu_exp_penal": self.loss_fn.cfg.mu_exp_penal,
+                    }
+
+                    logger.info("Synergy search: tentando %d configurações...", self.cfg.n_synergy_attempts)
+                    best_cfg = self.synergy.search(
+                        base_config=base_cfg,
+                        eval_fn=lambda c: self._eval_config(c, sample_batch),
+                    )
+
+                    # Aplica a melhor configuração
+                    for k, v in best_cfg.items():
+                        if hasattr(self.loss_fn.cfg, k):
+                            setattr(self.loss_fn.cfg, k, v)
+                    logger.info("Synergy search concluído. Melhor config aplicada.")
+                except Exception as e:
+                    logger.warning("Synergy search falhou: %s", e)
+
+            # Loop de batches
+            self.model.train()
+            if self.generator: self.generator.train()
+            if self.verifier: self.verifier.train()
+            if self.anti_hall: self.anti_hall.train()
+
+            for batch_idx, batch in enumerate(dataloader):
+                if batch_idx >= self.cfg.max_batches_per_epoch:
+                    break
+
+                try:
+                    metrics = self._train_step(batch)
+                    epoch_loss += metrics["loss"]
+                    epoch_ppl += metrics["ppl"]
+                    n_batches += 1
+                    if metrics["hypothesis_activated"]:
+                        hypothesis_activations += 1
+
+                    # Armazena métricas
+                    for k in ["loss", "ppl", "v_mean", "lr_G", "lr_V"]:
+                        self.metrics[k].append(metrics[k])
+
+                    if batch_idx % self.cfg.log_every == 0:
+                        logger.info(
+                            "Época %d | Batch %d | loss=%.4f | ppl=%.2f | v_mean=%.3f | lr_G=%.2e | hyp_act=%s",
+                            epoch + 1, batch_idx, metrics["loss"], metrics["ppl"],
+                            metrics["v_mean"], metrics["lr_G"],
+                            "Y" if metrics["hypothesis_activated"] else "N",
+                        )
+                except Exception as e:
+                    logger.error("Erro no batch %d: %s", batch_idx, e, exc_info=True)
+                    continue
+
+            # Cleanup de memória ao final da época
+            self.mem.cleanup()
+
+            avg_loss = epoch_loss / max(n_batches, 1)
+            avg_ppl = epoch_ppl / max(n_batches, 1)
+            logger.info(
+                "Época %d concluída | avg_loss=%.4f | avg_ppl=%.2f | hyp_activations=%d | batches_ok=%d/%d",
+                epoch + 1, avg_loss, avg_ppl, hypothesis_activations,
+                n_batches, self.cfg.max_batches_per_epoch,
+            )
+
+            # Detecta falha sistêmica: se nenhum batch teve sucesso, marca erro
+            if n_batches == 0:
+                logger.error("Época %d: NENHUM batch teve sucesso — possível bug no pipeline", epoch + 1)
+
+            self.history.append({
+                "epoch": epoch + 1,
+                "avg_loss": avg_loss,
+                "avg_ppl": avg_ppl,
+                "n_batches": n_batches,
+                "hypothesis_activations": hypothesis_activations,
+                "epoch_loss_total": epoch_loss,
+            })
+
+        elapsed = time.time() - start_time
+        logger.info("Treinamento concluído em %.1fs", elapsed)
+
+        return {
+            "history": self.history,
+            "metrics": self.metrics,
+            "final_loss": self.history[-1]["avg_loss"] if self.history else 0.0,
+            "final_ppl": self.history[-1]["avg_ppl"] if self.history else 0.0,
+            "elapsed_s": elapsed,
+            "synergy_history": self.synergy.history,
+            "auto_history": self.auto.history,
+        }
+
+
+__all__ = ["TrainerConfig", "CooperativeTrainer"]