rpotter6298
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
"""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)), {}
|