code_snapshots/WiSER_public_release_code/wiser/engine/trainer.py

"""Training driver for WiSER CIR set prediction.

`CsiSetTrainer.step(batch, gt)` wires together the cached scene-encoder,
the TX-modulator, the DETR head, the Hungarian matcher, and the loss bundle.

AC-5 invariant: the cache wrapper calls `model.scene_encode` at most `U` times
and `model.modulate_scene_tx` at most `P` times. Both invariants are tested
with U < P in `tests/test_cache_invariants.py`.
"""

from __future__ import annotations

import time as _time
from dataclasses import dataclass

import torch

from ..config import LossConfig as SharedLossConfig, ModelConfig as SharedModelConfig
from ..models.model import ModelConfig, SparseCsiDetrModel
from .cached_forward import CachedSceneTxForward
from .losses import LossWeights, csi_set_loss
from .matching import MatcherConfig, hungarian_match_batch


def _sync_now(is_cuda: bool) -> float:
    """Return `time.time()` after a CUDA sync when `is_cuda` is True.

    Consolidates the repeated `if _is_cuda: torch.cuda.synchronize(); t = time.time()`
    phase-timer boilerplate without changing timing semantics.
    """
    if is_cuda:
        torch.cuda.synchronize()
    return _time.time()


@dataclass
class TrainerConfig:
    model: ModelConfig = None  # type: ignore[assignment]
    loss_weights: LossWeights = None  # type: ignore[assignment]
    matcher: MatcherConfig = None  # type: ignore[assignment]

    @classmethod
    def from_shared(cls, shared_model: SharedModelConfig, shared_loss: SharedLossConfig) -> "TrainerConfig":
        return cls(
            model=ModelConfig.from_shared(shared_model),
            loss_weights=LossWeights.from_shared(shared_loss),
            matcher=MatcherConfig(),
        )


class CsiSetTrainer:
    """Minimal trainer exposing `step(batch, gt)` with split cache invariants."""

    def __init__(self, model: SparseCsiDetrModel, config: TrainerConfig | None = None) -> None:
        self.model = model
        self.cfg = config or TrainerConfig()
        if self.cfg.model is None:
            self.cfg.model = model.config
        if self.cfg.loss_weights is None:
            self.cfg.loss_weights = LossWeights()
        if self.cfg.matcher is None:
            self.cfg.matcher = MatcherConfig()
        self.cached_forward = CachedSceneTxForward()

    def _scene_encoder(self, voxel_level):
        """Called once per unique scene: runs the REAL backbone."""
        return self.model.scene_encode(voxel_level)

    def _tx_modulator(self, scene_tokens, tx_coord):
        """Called once per unique (scene, TX): cheap per-TX modulation."""
        tx_emb = self.model.encode_tx(tx_coord.unsqueeze(0))
        out = self.model.modulate_scene_tx(scene_tokens, tx_emb)
        return out["feats"]

    def _autocast_ctx(self):
        """bf16 autocast when on CUDA with the trellis2 backbone (FlashAttention
        requires fp16/bf16). Everything else runs fp32."""
        use_trellis = getattr(self.model.config.backbone, "kind", "") == "trellis2"
        device_type = "cuda" if torch.cuda.is_available() else "cpu"
        if use_trellis and device_type == "cuda":
            return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
        # No-op context manager.
        from contextlib import nullcontext
        return nullcontext()

    def step(self, batch: dict, gt: dict) -> dict:
        _is_cuda = torch.cuda.is_available()
        with self._autocast_ctx():
            tx_mem = self.cached_forward.forward_batch(
                batch,
                scene_encoder=self._scene_encoder,
                tx_modulator=self._tx_modulator,
            )
            if tx_mem.dim() == 4:
                tx_mem = tx_mem.squeeze(1)
            rx_xyz_norm = batch["rx_xyz_norm"]
            # Round-4 task20: time `decoder` (cross-attention) and `head`
            # (three output projections) separately, per Codex R3 finding 1.
            _t_dec = _sync_now(_is_cuda)
            decoded = self.model.head.decode(tx_mem, rx_xyz_norm)
            _t_proj = _sync_now(_is_cuda)
            decoder_s = _t_proj - _t_dec
            predictions = self.model.head.project(decoded)
            head_s = _sync_now(_is_cuda) - _t_proj

        # Round-3 task20 contract: time `hungarian_match_batch` and
        # `csi_set_loss` separately so the phase table covers the ENTIRE
        # time spent inside `driver.step(...)` (Codex R2 finding 2).
        _t_match = _sync_now(_is_cuda)
        matchings = hungarian_match_batch(predictions, gt, self.cfg.matcher)
        _t_loss = _sync_now(_is_cuda)
        match_s = _t_loss - _t_match

        loss_bundle = csi_set_loss(predictions, gt, matchings, weights=self.cfg.loss_weights)
        loss_s = _sync_now(_is_cuda) - _t_loss

        return {
            "predictions": predictions,
            "matchings": matchings,
            "loss_bundle": loss_bundle,
            "cache_counters": {
                "c_enc": self.cached_forward.counters.c_enc,
                "c_txmod": self.cached_forward.counters.c_txmod,
                "u": self.cached_forward.counters.u,
                "p": self.cached_forward.counters.p,
                "q": self.cached_forward.counters.q,
            },
            # Round-3 phase timers (seconds on this step). `decoder_plus_head`
            # remains merged because `model.head_forward` wraps the cross-
            # attention decoder and the three output projections as one
            # compiled nn.Module graph. `matching` and `loss_construction`
            # are new buckets added in Round 3 so the phase table accounts
            # for the entire `driver.step(...)` wall-clock time. The outer
            # trainer (`scripts/train_and_eval.py`) additionally records
            # `dataloader`, `backward`, and `optimizer_step`.
            "phase_times_s": {
                "scene_encode": float(self.cached_forward.counters.scene_encode_s),
                "tx_modulate": float(self.cached_forward.counters.tx_modulate_s),
                "decoder": float(decoder_s),
                "head": float(head_s),
                "matching": float(match_s),
                "loss_construction": float(loss_s),
            },
        }


__all__ = ["TrainerConfig", "CsiSetTrainer"]