v4 update

v4/classes/vehicles-deprecated/htbase.py

@@ -0,0 +1,40 @@
+"""htbase — HTBase: abstract base for all v4 vehicle classes."""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+
+class HTBase(nn.Module, ABC):
+    """Shared interface for all HyperTower vehicles.
+
+    Subclasses must implement ``encode`` and ``forward``.
+
+    ``transform`` walks ``self.towers`` (if present) and returns the transform
+    from the first tower that exposes one — used by data loaders.
+
+    ``forward`` contract: returns ``(logits, aux_dict)`` where
+    ``aux_dict`` maps a name or index to per-component logits.
+    HTMono returns an empty dict to keep the signature uniform.
+    """
+
+    @abstractmethod
+    def encode(self, inputs) -> torch.Tensor:
+        """Return the pre-classifier embedding."""
+
+    @abstractmethod
+    def forward(self, inputs) -> tuple[torch.Tensor, dict]:
+        """Return (logits, aux_dict)."""
+
+    @property
+    def transform(self):
+        towers = getattr(self, "towers", None) or {}
+        for t in (towers.values() if hasattr(towers, "values") else []):
+            if hasattr(t, "transform"):
+                return t.transform
+        encoder = getattr(self, "encoder", None)
+        if encoder is not None:
+            return getattr(encoder, "transform", None)
+        return None

v4/classes/vehicles-deprecated/htfusion.py

@@ -0,0 +1,61 @@
+"""htfusion — HTFusion: N named towers fused through a FusionBridge."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.bridges.fusion_bridge import FusionBridge
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTFusion(HTBase):
+    """General N-tower fusion vehicle.
+
+    Each named encoder is registered as a submodule; the FusionBridge
+    projects and Hadamard-fuses their embeddings.
+
+    Parameters
+    ----------
+    towers      : ordered dict ``{name: encoder}``.  Each encoder must
+                  expose ``.out_dim``.
+    num_classes : output classes
+    fusion_dim  : bridge projection dimensionality
+    dropout     : bridge dropout
+    use_se      : SE gate on the fused vector
+
+    Forward contract
+    ----------------
+    ``forward(embeddings)`` takes a ``dict[str, Tensor]`` of pre-computed
+    per-tower embeddings and returns ``(logits_fused, aux_dict)`` where
+    ``aux_dict`` maps each tower name to its auxiliary head logits.
+    """
+
+    def __init__(
+        self,
+        towers:      dict[str, nn.Module],
+        num_classes: int,
+        fusion_dim:  int   = 256,
+        dropout:     float = 0.5,
+        use_se:      bool  = False,
+    ):
+        super().__init__()
+        self.towers = nn.ModuleDict(towers)
+        self.bridge = FusionBridge(
+            tower_dims=[t.out_dim for t in self.towers.values()],
+            num_classes=num_classes,
+            fusion_dim=fusion_dim,
+            dropout=dropout,
+            use_se=use_se,
+        )
+
+    def encode(self, embeddings: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return z_fused (pre-classifier) from a dict of per-tower embeddings."""
+        return self.bridge.encode([embeddings[name] for name in self.towers])
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        ordered = [embeddings[name] for name in self.towers]
+        logits, aux = self.bridge.fuse(ordered)
+        return logits, {name: aux[i] for i, name in enumerate(self.towers)}

v4/classes/vehicles-deprecated/htlateral.py

@@ -0,0 +1,67 @@
+"""htlateral — HTLateral: shared encoder over N same-type inputs."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTLateral(HTBase):
+    """N same-type inputs through a shared encoder, jointly compressed, then classified.
+
+    All inputs share the same encoder weights (one forward pass per input).
+    The joint MLP compresses the concatenated embeddings before classification.
+
+    Aux heads provide per-input logits before the joint MLP — useful for
+    BCD-style training.
+
+    Parameters
+    ----------
+    encoder     : shared encoder module with ``.out_dim``
+    input_names : ordered slot names (e.g. ``["od", "os"]``)
+    num_classes : output classes
+    fusion_dim  : joint MLP hidden dim
+    dropout     : dropout in MLP and classifier
+    """
+
+    def __init__(
+        self,
+        encoder:     nn.Module,
+        input_names: list[str],
+        num_classes: int,
+        fusion_dim:  int   = 256,
+        dropout:     float = 0.5,
+    ):
+        super().__init__()
+        self.encoder     = encoder
+        self.input_names = list(input_names)
+        n       = len(input_names)
+        in_dim: int = encoder.out_dim  # type: ignore[assignment]
+
+        self.joint = nn.Sequential(
+            nn.Linear(n * in_dim, fusion_dim), nn.LayerNorm(fusion_dim),
+            nn.ReLU(), nn.Dropout(dropout), nn.Linear(fusion_dim, in_dim),
+        )
+        self.aux_heads = nn.ModuleList([
+            nn.Linear(in_dim, num_classes) for _ in range(n)
+        ])
+        self.head = nn.Sequential(
+            nn.ReLU(), nn.Dropout(dropout), nn.Linear(in_dim, num_classes),
+        )
+
+    def encode(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return joint embedding (post-MLP, pre-classifier)."""
+        zs = [self.encoder(inputs[name]) for name in self.input_names]
+        return self.joint(torch.cat(zs, dim=1))
+
+    def forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        zs      = [self.encoder(inputs[name]) for name in self.input_names]
+        z_joint = self.joint(torch.cat(zs, dim=1))
+        logits  = self.head(z_joint)
+        aux     = {name: head(z)
+                   for name, head, z in zip(self.input_names, self.aux_heads, zs)}
+        return logits, aux

v4/classes/vehicles-deprecated/htmono.py

@@ -0,0 +1,44 @@
+"""htmono — HTMono: single tower + ClassificationHead, no bridge."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.heads.classifier import ClassificationHead
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTMono(HTBase):
+    """Single-tower vehicle: tower embedding fed directly into a ClassificationHead.
+
+    No bridge or projection — the tower's output goes straight to
+    ReLU → Dropout → Linear.  Returns ``(logits, {})`` from ``forward``
+    to match the HTFusion / HTLateral interface.
+
+    Parameters
+    ----------
+    tower       : encoder module with ``.out_dim``
+    num_classes : output classes
+    dropout     : dropout before the output linear layer
+    """
+
+    def __init__(
+        self,
+        tower:       nn.Module,
+        num_classes: int,
+        dropout:     float = 0.5,
+    ):
+        super().__init__()
+        self.tower = tower
+        self.head  = ClassificationHead(tower.out_dim, num_classes, dropout)  # type: ignore[arg-type]
+
+    def encode(self, inputs) -> torch.Tensor:
+        """Return tower embedding (pre-classifier)."""
+        if isinstance(inputs, dict):
+            # single-entry dict from HTDataset eye-level pass
+            (z,) = inputs.values()
+            return self.tower(z) if torch.is_tensor(z) else self.tower(*z.values())
+        return self.tower(inputs)
+
+    def forward(self, inputs) -> tuple[torch.Tensor, dict]:
+        return self.head(self.encode(inputs)), {}