| |
| """Experiment 8-v2: stronger unified RGB+pose AfriSign encoder. |
| |
| This file reuses the Exp8 data pipeline and training loop, but swaps in a |
| stronger shared encoder: |
| |
| - temporal convolutional pose stem before the Transformer |
| - language/modality/task-conditioned residual adapters |
| - stronger projection head for supervised contrastive learning |
| |
| It is intentionally a thin wrapper around `exp8_unified_mixed_encoder.py` so the |
| same manifests, loaders, metrics, and aggregation scripts remain compatible. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| from experiments import exp8_unified_mixed_encoder as exp8 |
|
|
|
|
| class TemporalPoseStem(nn.Module): |
| """TransSLR-style temporal stem for landmark sequences.""" |
|
|
| def __init__(self, feature_dim: int, hidden_dim: int, dropout: float) -> None: |
| super().__init__() |
| self.in_proj = nn.Linear(feature_dim, hidden_dim) |
| self.blocks = nn.ModuleList( |
| [ |
| nn.Sequential( |
| nn.LayerNorm(hidden_dim), |
| nn.Conv1d(hidden_dim, hidden_dim, kernel_size=5, padding=2, groups=hidden_dim), |
| nn.GELU(), |
| nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1), |
| nn.Dropout(dropout), |
| ) |
| for _ in range(3) |
| ] |
| ) |
| self.out_norm = nn.LayerNorm(hidden_dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.in_proj(x) |
| for block in self.blocks: |
| y = block[0](x).transpose(1, 2) |
| y = block[1](y) |
| y = block[2](y) |
| y = block[3](y).transpose(1, 2) |
| y = block[4](y) |
| x = x + y |
| return self.out_norm(x) |
|
|
|
|
| class ConditionalAdapter(nn.Module): |
| """Small residual adapter conditioned by language/modality/level/task context.""" |
|
|
| def __init__(self, hidden_dim: int, bottleneck: int, dropout: float) -> None: |
| super().__init__() |
| self.norm = nn.LayerNorm(hidden_dim) |
| self.down = nn.Linear(hidden_dim, bottleneck) |
| self.up = nn.Linear(bottleneck, hidden_dim) |
| self.drop = nn.Dropout(dropout) |
| self.gate = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.Sigmoid()) |
|
|
| def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: |
| update = self.up(torch.nn.functional.gelu(self.down(self.norm(x)))) |
| gate = self.gate(context).unsqueeze(1) |
| return x + self.drop(update * gate) |
|
|
|
|
| class StrongUnifiedAfriSignEncoder(nn.Module): |
| def __init__( |
| self, |
| *, |
| task_dims: dict[str, int], |
| num_languages: int, |
| hidden_dim: int, |
| feature_dim: int, |
| max_tokens: int, |
| layers: int, |
| heads: int, |
| ff_dim: int, |
| dropout: float, |
| rgb_backbone: str, |
| rgb_train_backbone: str, |
| rgb_pretrained: bool, |
| ) -> None: |
| super().__init__() |
| self.task_keys = list(task_dims) |
| self.pose_stem = TemporalPoseStem(feature_dim, hidden_dim, dropout) |
| if rgb_backbone == "small_cnn": |
| self.rgb_cnn = exp8.SmallFrameCNN(hidden_dim) |
| elif rgb_backbone == "efficientnet_b0": |
| self.rgb_cnn = exp8.EfficientNetB0FrameEncoder(hidden_dim, rgb_train_backbone, rgb_pretrained) |
| else: |
| raise ValueError(f"Unknown rgb_backbone={rgb_backbone!r}") |
|
|
| self.cls = nn.Parameter(torch.zeros(1, 1, hidden_dim)) |
| self.pos = nn.Embedding(max_tokens + 1, hidden_dim) |
| self.lang_emb = nn.Embedding(num_languages, hidden_dim) |
| self.modality_emb = nn.Embedding(len(exp8.MODALITY_IDS), hidden_dim) |
| self.level_emb = nn.Embedding(len(exp8.LEVEL_IDS), hidden_dim) |
| self.task_emb = nn.Embedding(len(self.task_keys), hidden_dim) |
|
|
| self.pre_adapter = ConditionalAdapter(hidden_dim, max(32, hidden_dim // 4), dropout) |
| enc_layer = nn.TransformerEncoderLayer( |
| d_model=hidden_dim, |
| nhead=heads, |
| dim_feedforward=ff_dim, |
| dropout=dropout, |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.encoder = nn.TransformerEncoder(enc_layer, layers, enable_nested_tensor=False) |
| self.post_adapter = ConditionalAdapter(hidden_dim, max(32, hidden_dim // 4), dropout) |
| self.norm = nn.LayerNorm(hidden_dim) |
| self.drop = nn.Dropout(dropout) |
| self.heads = nn.ModuleDict({key: nn.Linear(hidden_dim, dim) for key, dim in task_dims.items()}) |
| self.projector = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.GELU(), |
| nn.LayerNorm(hidden_dim), |
| nn.Linear(hidden_dim, hidden_dim), |
| ) |
| nn.init.trunc_normal_(self.cls, std=0.02) |
|
|
| def encode(self, batch: dict[str, torch.Tensor]) -> torch.Tensor: |
| if "pose" in batch: |
| x = self.pose_stem(batch["pose"]) |
| elif "rgb" in batch: |
| rgb = batch["rgb"] |
| b, t, c, h, w = rgb.shape |
| x = self.rgb_cnn(rgb.reshape(b * t, c, h, w)).reshape(b, t, -1) |
| else: |
| raise ValueError("Batch must contain pose or rgb") |
|
|
| b, t, _ = x.shape |
| positions = torch.arange(t + 1, device=x.device) |
| context = ( |
| self.lang_emb(batch["lang_idx"]) |
| + self.modality_emb(batch["modality_idx"]) |
| + self.level_emb(batch["level_idx"]) |
| + self.task_emb(batch["task_idx"]) |
| ) |
| x = torch.cat([self.cls.expand(b, -1, -1), x], dim=1) |
| x = x + self.pos(positions).unsqueeze(0) + context.unsqueeze(1) |
| x = self.pre_adapter(x, context) |
| x = self.encoder(x) |
| x = self.post_adapter(x, context) |
| return self.drop(self.norm(x[:, 0])) |
|
|
| def forward(self, batch: dict[str, torch.Tensor], task_key: str) -> torch.Tensor: |
| return self.heads[task_key](self.encode(batch)) |
|
|
| def contrast_features(self, features: torch.Tensor) -> torch.Tensor: |
| return torch.nn.functional.normalize(self.projector(features), dim=1) |
|
|
|
|
| def main() -> None: |
| exp8.UnifiedAfriSignEncoder = StrongUnifiedAfriSignEncoder |
| exp8.main() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
