Training Module¶

def train_epoch(
    self, progress: Progress | None = None, task_id: TaskID | None = None
) -> tuple[float, float]:
    """
    Train for one epoch.

    Args:
        progress: Rich Progress instance (optional)
        task_id: Progress task ID (optional)

    Returns:
        Tuple of (average_loss, accuracy)
    """
    self.model.train()
    if self.adversarial_head is not None:
        self.adversarial_head.train()
    if self.cl_regression_head is not None:
        self.cl_regression_head.train()

    # Training metrics reduce to a mean loss and an accuracy, so nothing
    # here needs the predictions themselves -- a running correct-count is
    # enough and keeps the per-epoch prediction arrays (one boxed
    # np.float32 per prediction) from ever existing. Totals stay on the
    # device until the epoch ends; see _DeviceTally.
    tally = _DeviceTally()
    seen = 0

    for batch in self.train_loader:
        self.optimizer.zero_grad()

        # Gradient accumulation: N sub-batches, each scaled by 1/N, one
        # optimizer step per full batch. N=1 leaves the graph untouched.
        sub_batches = split_batch(batch, self.grad_accum_split)
        num_splits = len(sub_batches)

        for sub_batch in sub_batches:
            if self.use_mixed_precision:
                with torch.amp.autocast("cuda"):
                    logits, labels, main_loss, loss, adv, cl_loss = self._compute_batch_loss(
                        sub_batch
                    )
                if num_splits > 1:
                    loss = loss / num_splits
                self.scaler.scale(loss).backward()
            else:
                logits, labels, main_loss, loss, adv, cl_loss = self._compute_batch_loss(
                    sub_batch
                )
                if num_splits > 1:
                    loss = loss / num_splits
                loss.backward()

            # Track metrics (sub-batch losses are averaged back to a
            # per-batch mean so history is comparable across split settings)
            tally.add("loss", main_loss, num_splits)
            if adv is not None:
                adv_loss, adv_preds, adv_labels = adv
                tally.add("adv_loss", adv_loss, num_splits)
                # `mask.any()` would sync; the masked matches and the
                # masked count are all the accuracy needs, and both are
                # correct when the mask is empty.
                mask = adv_labels != -1
                tally.add("adv_correct", ((adv_preds == adv_labels) & mask).sum())
                tally.add("adv_seen", mask.sum())
            if cl_loss is not None:
                tally.add("cl_loss", cl_loss, num_splits)

            labels_flat = labels.detach().flatten()
            if self.loss_type == "cross_entropy" and self._num_out > 2:
                # Multi-class: argmax predictions
                preds = torch.argmax(logits, dim=-1).detach().flatten()
            else:
                if self.loss_type == "cross_entropy":
                    # Binary CE: probabilities via softmax, take class 1
                    scores = torch.softmax(logits, dim=-1)[:, 1]
                else:
                    scores = torch.sigmoid(logits)
                preds = scores.detach().flatten() > 0.5
            matches = preds.to(torch.float64) == labels_flat.to(torch.float64)
            tally.add("correct", matches.sum())
            seen += labels_flat.numel()

        # Gradient clipping + optimizer step (once per full batch)
        needs_clip = self.clip_grad_fn is not None or self.max_grad_norm > 0
        if needs_clip:
            if self.use_mixed_precision:
                self.scaler.unscale_(self.optimizer)
            if self.clip_grad_fn is not None:
                self.clip_grad_fn(self.model.parameters())
            else:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

        if self.use_mixed_precision:
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()

        # Update progress if provided
        if progress is not None and task_id is not None:
            progress.update(task_id, advance=1)

    # Compute metrics. accuracy_score normalizes to matches / samples, so
    # the running counts reproduce it exactly for both the multi-class
    # (argmax) and the binary (threshold at 0.5) case.
    avg_loss = tally.value("loss") / len(self.train_loader)
    accuracy = tally.value("correct") / seen if seen else 0.0

    # Adversarial metrics
    if self.adversarial_head is not None:
        avg_adv_loss = tally.value("adv_loss") / max(1, len(self.train_loader))
        adv_seen = tally.value("adv_seen")
        adv_acc = tally.value("adv_correct") / adv_seen if adv_seen else 0.0
        self.history["train_adv_loss"].append(avg_adv_loss)
        self.history["train_adv_acc"].append(adv_acc)

    # CL regression metrics
    if self.cl_regression_head is not None:
        avg_cl_loss = tally.value("cl_loss") / max(1, len(self.train_loader))
        self.history["train_cl_loss"].append(avg_cl_loss)

    return avg_loss, accuracy

def validate(
    self, progress: Progress | None = None, task_id: TaskID | None = None
) -> tuple[float, float, float, float]:
    """
    Validate model.

    Args:
        progress: Rich Progress instance (optional)
        task_id: Progress task ID (optional)

    Returns:
        Tuple of (average_loss, accuracy, roc_auc, f1)
    """
    if self.val_loader is None:
        return 0.0, 0.0, 0.0, 0.0

    self.model.eval()
    if self.cl_regression_head is not None:
        self.cl_regression_head.eval()
    # AUROC needs the probabilities themselves, so validation keeps them --
    # but as one array per batch, never as a list of boxed numpy scalars
    # (~40 bytes each while alive, and unboxed one at a time at the end).
    tally = _DeviceTally()
    all_preds: list[np.ndarray] = []
    all_labels: list[np.ndarray] = []
    # For multiclass: keep full softmax probabilities so we can compute
    # macro one-vs-rest AUROC (argmax-only preds lose that information).
    all_probs_mc: list[np.ndarray] = []

    with torch.inference_mode():
        for batch in self.val_loader:
            # Move labels to device
            labels = batch["label"].to(self.device)

            # Adapt labels for CrossEntropyLoss
            if self.loss_type == "cross_entropy":
                ce_labels = labels.squeeze(-1).long()
            else:
                ce_labels = None

            # Forward pass (wrapper handles moving tensors and calling model correctly)
            if self.use_mixed_precision:
                with torch.amp.autocast("cuda"):
                    logits = self.model_wrapper.forward_batch(batch, self.device)
                    if ce_labels is not None:
                        loss = self.criterion(logits, ce_labels)
                    else:
                        loss = self.criterion(logits, labels)
            else:
                logits = self.model_wrapper.forward_batch(batch, self.device)
                if ce_labels is not None:
                    loss = self.criterion(logits, ce_labels)
                else:
                    loss = self.criterion(logits, labels)

            # CL regression validation loss
            if (
                self.cl_regression_head is not None
                and "cl_target" in batch
                and self.model_wrapper.captured_repr is not None
            ):
                cl_targets = batch["cl_target"].to(self.device)
                cl_mask = cl_targets >= 0
                if cl_mask.any():
                    cl_preds = self.cl_regression_head(
                        self.model_wrapper.captured_repr[cl_mask]
                    )
                    tally.add(
                        "cl_loss",
                        nn.functional.mse_loss(cl_preds, cl_targets[cl_mask]),
                    )

            # Track metrics
            tally.add("loss", loss)
            if self.loss_type == "cross_entropy" and self._num_out > 2:
                probs_mc = torch.softmax(logits, dim=-1).cpu().numpy()
                all_probs_mc.append(probs_mc)
                all_preds.append(probs_mc.argmax(axis=-1).ravel())
            elif self.loss_type == "cross_entropy":
                all_preds.append(torch.softmax(logits, dim=-1)[:, 1].cpu().numpy().ravel())
            else:
                all_preds.append(torch.sigmoid(logits).cpu().numpy().ravel())
            all_labels.append(labels.cpu().numpy().ravel())

            # Update progress if provided
            if progress is not None and task_id is not None:
                progress.update(task_id, advance=1)

    # Compute metrics
    avg_loss = tally.value("loss") / len(self.val_loader)
    preds_arr = np.concatenate(all_preds) if all_preds else np.zeros(0, dtype=np.float32)
    labels_arr = np.concatenate(all_labels) if all_labels else np.zeros(0, dtype=np.float32)
    if self._num_out > 2:
        # Multi-class: preds are class indices; probs are needed for AUROC.
        accuracy = accuracy_score(labels_arr, preds_arr)
        f1 = f1_score(labels_arr, preds_arr, average="macro", zero_division=0.0)
        present = np.unique(labels_arr)
        # Macro one-vs-rest AUROC computed per-class. We avoid passing the
        # sliced probs to roc_auc_score(multi_class="ovr") because sklearn
        # rejects scores that don't sum to 1.0 across columns — which fires
        # whenever the val set is missing any of the model's output classes
        # (silent ValueError → AUC reported as 0.0). Manual OvR sidesteps
        # the check and lets us skip classes that are degenerate (zero
        # positives or zero negatives) in this validation set.
        if len(present) > 1 and all_probs_mc:
            probs_all = np.concatenate(all_probs_mc, axis=0)
            per_class_aucs: list[float] = []
            for c in present:
                c_int = int(c)
                if c_int >= probs_all.shape[1]:
                    continue
                y_bin = (labels_arr == c).astype(int)
                pos = int(y_bin.sum())
                if pos == 0 or pos == len(y_bin):
                    continue
                try:
                    per_class_aucs.append(float(roc_auc_score(y_bin, probs_all[:, c_int])))
                except ValueError as e:
                    logger.warning(f"AUROC failed for class {c_int}: {e}")
            auc = float(np.mean(per_class_aucs)) if per_class_aucs else 0.0
        else:
            auc = 0.0
    else:
        all_preds_binary = (preds_arr > 0.5).astype(int)
        accuracy = accuracy_score(labels_arr, all_preds_binary)
        auc = roc_auc_score(labels_arr, preds_arr) if len(np.unique(labels_arr)) > 1 else 0.0
        f1 = f1_score(labels_arr, all_preds_binary, zero_division=0.0)

    # CL regression validation metrics
    if self.cl_regression_head is not None:
        avg_cl_loss = tally.value("cl_loss") / max(1, len(self.val_loader))
        self.history["val_cl_loss"].append(avg_cl_loss)

    return avg_loss, accuracy, auc, f1

def train(self, epochs: int, early_stopping_patience: int = 10) -> dict[str, Any]:
    """
    Train model for multiple epochs.

    Args:
        epochs: Number of epochs to train
        early_stopping_patience: Stop if validation accuracy doesn't improve for N epochs (0 to disable)

    Returns:
        Training history dictionary
    """
    patience_counter = 0
    total_epochs = max(0, epochs - self.start_epoch + 1)
    last_epoch = self.start_epoch - 1

    # If training already completed (resume past final epoch), save best and exit
    if self.start_epoch > epochs:
        logger.info(
            f"Training already complete (resumed at epoch {self.start_epoch - 1}, "
            f"requested {epochs}). Saving checkpoints and exiting."
        )
        if self.output_dir:
            self._ensure_best_checkpoint()
            self.save_checkpoint("model_last.pt", epoch=self.start_epoch - 1)
            self.save_history()
        return self.history

    # Create progress bars
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
        console=console,
    ) as progress:
        epoch_task = progress.add_task("[cyan]Training epochs...", total=total_epochs)

        for epoch in range(self.start_epoch, epochs + 1):
            last_epoch = epoch
            # LR warmup. The cosine schedule folds warmup into its own
            # LambdaLR multiplier, so only the plateau/none paths need the
            # manual ramp here.
            if (
                self.scheduler_type != "cosine"
                and self.warmup_epochs > 0
                and epoch <= self.warmup_epochs
            ):
                warmup_lr = self.base_lr * (epoch / self.warmup_epochs)
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = warmup_lr

            # Adversarial lambda annealing (linear ramp)
            if self.adversarial_head is not None and self.adversarial_anneal_epochs > 0:
                progress_frac = min(1.0, epoch / self.adversarial_anneal_epochs)
                current_lambda = self.adversarial_lambda * progress_frac
                self.adversarial_head.set_lambda(current_lambda)

            # Create tasks for this epoch
            train_task = progress.add_task(
                f"[green]Epoch {epoch}/{epochs} - Training", total=len(self.train_loader)
            )

            # Train
            train_loss, train_acc = self.train_epoch(progress, train_task)
            self.history["train_loss"].append(train_loss)
            self.history["train_acc"].append(train_acc)

            progress.remove_task(train_task)

            # Validate
            if self.val_loader is not None:
                val_task = progress.add_task(
                    f"[yellow]Epoch {epoch}/{epochs} - Validation", total=len(self.val_loader)
                )
                val_loss, val_acc, val_auc, val_f1 = self.validate(progress, val_task)
                self.history["val_loss"].append(val_loss)
                self.history["val_acc"].append(val_acc)
                self.history["val_auc"].append(val_auc)
                self.history["val_f1"].append(val_f1)

                progress.remove_task(val_task)

                # LR scheduler step (only after warmup). For reduce_on_plateau
                # we step on the *selection metric*, not val_loss, so the
                # plateau detection matches what early-stopping and best-
                # model checkpointing react to. If val_loss were used while
                # selection runs off F1/AUC, LR can drop because loss
                # plateaued on easy examples even though the ranking metric
                # is still improving — burning the schedule prematurely.
                # The cosine LambdaLR owns warmup itself, so it steps every
                # epoch; plateau keeps the original post-warmup guard.
                if self.scheduler is not None and (
                    self.scheduler_type == "cosine" or epoch > self.warmup_epochs
                ):
                    old_lr = self.optimizer.param_groups[0]["lr"]
                    if self.scheduler_type == "reduce_on_plateau":
                        if self._checkpoint_metric == "val_f1":
                            plateau_signal = -val_f1
                        elif self._checkpoint_metric == "val_auc":
                            plateau_signal = -val_auc
                        elif self._checkpoint_metric == "val_acc":
                            plateau_signal = -val_acc
                        else:
                            plateau_signal = val_loss
                        self.scheduler.step(plateau_signal)
                    else:
                        # Epoch-based schedulers (cosine, etc.)
                        self.scheduler.step()
                    new_lr = self.optimizer.param_groups[0]["lr"]
                    if new_lr != old_lr:
                        logger.info(f"LR reduced: {old_lr:.6f} -> {new_lr:.6f}")

                # Display metrics — [*] marks the checkpoint criterion
                lr_str = ""
                current_lr = self.optimizer.param_groups[0]["lr"]
                if current_lr != self.base_lr:
                    lr_str = f" LR: {current_lr:.6f}"
                acc_label = "Acc[*]" if self._checkpoint_metric == "val_acc" else "Acc"
                f1_label = "F1[*]" if self._checkpoint_metric == "val_f1" else "F1"
                auc_label = "AUC[*]" if self._checkpoint_metric == "val_auc" else "AUC"
                adv_str = ""
                if self.adversarial_head is not None and self.history["train_adv_loss"]:
                    _adv_l = self.history["train_adv_loss"][-1]
                    _adv_a = self.history["train_adv_acc"][-1]
                    _lam = self.adversarial_head.grl.lambda_
                    adv_str = f" | Adv: L={_adv_l:.3f} Acc={_adv_a:.3f} λ={_lam:.3f}"
                cl_str = ""
                if self.cl_regression_head is not None and self.history["train_cl_loss"]:
                    _tcl = self.history["train_cl_loss"][-1]
                    _vcl = (
                        self.history["val_cl_loss"][-1] if self.history["val_cl_loss"] else 0.0
                    )
                    cl_str = f" | CL: train={_tcl:.4f} val={_vcl:.4f}"
                console.print(
                    f"[cyan]Epoch {epoch}/{epochs}[/cyan] | "
                    f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
                    f"Val Loss: {val_loss:.4f} {acc_label}: {val_acc:.4f} "
                    f"{f1_label}: {val_f1:.4f} {auc_label}: {val_auc:.4f}"
                    f"{lr_str}{adv_str}{cl_str}"
                )

                # Save best model per configured checkpoint metric
                if self._checkpoint_metric == "val_f1":
                    improved = val_f1 > self.best_val_f1
                elif self._checkpoint_metric == "val_auc":
                    improved = val_auc > self.best_val_auc
                else:
                    improved = val_acc > self.best_val_acc
                if improved:
                    self.best_val_acc = val_acc
                    self.best_val_f1 = val_f1
                    self.best_val_auc = val_auc
                    self.best_epoch = epoch
                    self._best_model_state = copy.deepcopy(self.model.state_dict())
                    patience_counter = 0

                    if self.output_dir:
                        self.save_checkpoint("model_best.pt", epoch=epoch)
                        console.print(
                            f"[bold green]✓ Saved best model "
                            f"(val_acc: {val_acc:.4f}, val_f1: {val_f1:.4f}, "
                            f"val_auc: {val_auc:.4f})[/bold green]"
                        )
                else:
                    patience_counter += 1

                # Early stopping (disabled if patience is 0)
                if early_stopping_patience > 0 and patience_counter >= early_stopping_patience:
                    console.print(f"[yellow]Early stopping at epoch {epoch}[/yellow]")
                    break
            else:
                console.print(
                    f"[cyan]Epoch {epoch}/{epochs}[/cyan] | "
                    f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f}"
                )

            progress.update(epoch_task, advance=1)

    # Save final model
    if self.output_dir:
        self.save_checkpoint("model_last.pt", epoch=last_epoch)
        self._ensure_best_checkpoint()
        self.save_history()

    return self.history

def save_checkpoint(self, filename: str, epoch: int = 0) -> None:
    """Save model checkpoint."""
    if self.output_dir is None:
        return

    checkpoint_path = self.output_dir / filename
    checkpoint: dict[str, Any] = {
        "model_state_dict": self.model.state_dict(),
        "best_val_acc": self.best_val_acc,
        "best_val_f1": self.best_val_f1,
        "best_val_auc": self.best_val_auc,
        "best_epoch": self.best_epoch,
        "epoch": epoch,
        "scheduler_state_dict": self.scheduler.state_dict() if self.scheduler else None,
        "scaler_state_dict": self.scaler.state_dict() if self.scaler else None,
        "best_model_state_dict": self._best_model_state,
    }
    if self._should_save_optimizer(epoch):
        checkpoint["optimizer_state_dict"] = self.optimizer.state_dict()
    if self.cl_regression_head is not None:
        checkpoint["cl_regression_head_state_dict"] = self.cl_regression_head.state_dict()
    torch.save(checkpoint, checkpoint_path)

def save_history(self) -> None:
    """Save training history to JSON."""
    if self.output_dir is None:
        return

    history_path = self.output_dir / "metrics.json"
    with open(history_path, "w") as f:
        json.dump(self.history, f, indent=2)

    # Save summary (only if training occurred)
    if len(self.history["train_loss"]) == 0:
        return

    summary = {
        "best_val_acc": self.best_val_acc,
        "best_val_f1": self.best_val_f1,
        "best_val_auc": self.best_val_auc,
        "best_epoch": self.best_epoch,
        "checkpoint_metric": self._checkpoint_metric,
        "final_train_loss": self.history["train_loss"][-1],
        "final_train_acc": self.history["train_acc"][-1],
    }

    if self.val_loader is not None and len(self.history["val_loss"]) > 0:
        summary.update(
            {
                "final_val_loss": self.history["val_loss"][-1],
                "final_val_acc": self.history["val_acc"][-1],
                "final_val_auc": self.history["val_auc"][-1],
                "final_val_f1": self.history["val_f1"][-1],
            }
        )

    summary_path = self.output_dir / "summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)