rpotter6298
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
"""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
|