experiment-2107/scripts/model_transformer_bilstm.py

"""Transformer-encoder -> BiLSTM-decoder architecture for figure-skating action classification.

Two temporal-modeling stages, no conv backbone:
  1. Transformer encoder: linear input projection + sinusoidal positional encoding + a stack
     of multi-head self-attention blocks (global context across all timesteps).
  2. BiLSTM decoder: a bidirectional LSTM that sweeps the encoder's output sequentially,
     re-imposing local temporal order before pooling.
Followed by temporal mean pooling and the same deep GELU dense head used in model.py, so the
classifier itself is unchanged and results are comparable to the conv-based variants.
Trained under the same settings (LR, batch size, class weighting, early stopping) as model.py.

Usage:
    python model_transformer_bilstm.py --data-dir /path/to/processed
    python model_transformer_bilstm.py --data-dir /path/to/processed --coarse
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import f1_score
from sklearn.utils.class_weight import compute_class_weight
from torch.utils.data import DataLoader, TensorDataset

try:
    from . import labels as labels_mod
except ImportError:
    import labels as labels_mod

from model import (
    DenseBlock,
    load_split,
    resolve_taxonomy,
    assert_label_consistency,
    report_metrics,
    run_epoch,
    predict,
)
from model_transformers import PositionalEncoding, make_transformer_encoder

# ---------------------------------------------------------------------------
# Config (identical to model.py so the comparison isolates the temporal-modeling stages)
# ---------------------------------------------------------------------------

EPOCHS = 100
BATCH_SIZE = 64
LEARNING_RATE = 5e-4
WEIGHT_DECAY = 1e-4
GRAD_CLIP = 1.0
MAX_CLASS_WEIGHT = 5.0
EARLY_STOP_PATIENCE = 20
SEED = 42
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

D_MODEL = 384
N_ATTN_BLOCKS = 3
LSTM_LAYERS = 2
LSTM_DROPOUT = 0.2


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------


class SkatingTransformerBiLSTMClassifier(nn.Module):
    """(B, T, F) sequence of skeleton features -> (B, num_classes) class logits.

    Stage 1 (encoder): linear projection to d_model + positional encoding + a real
    Transformer encoder (self-attn + FFN per layer) -- no conv backbone.
    Stage 2 (decoder): bidirectional LSTM over the encoder's output sequence.
    Then temporal mean pool -> the same dense head as model.py.
    """

    def __init__(self, in_features: int, num_classes: int, num_layers: int = N_ATTN_BLOCKS):
        super().__init__()
        self.input_proj = nn.Linear(in_features, D_MODEL)
        self.pos_enc = PositionalEncoding(D_MODEL)
        self.encoder = make_transformer_encoder(dim=D_MODEL, heads=8, num_layers=num_layers)
        self.decoder = nn.LSTM(
            D_MODEL, D_MODEL // 2, num_layers=LSTM_LAYERS, batch_first=True,
            bidirectional=True, dropout=LSTM_DROPOUT,
        )
        self.head = nn.Sequential(
            DenseBlock(D_MODEL, 1024, 0.5),
            DenseBlock(1024, 512, 0.4),
            DenseBlock(512, 256, 0.3),
            nn.LayerNorm(256),
        )
        self.out = nn.Linear(256, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)                      # (B, T, D_MODEL)
        x = self.pos_enc(x)
        x = self.encoder(x)                          # transformer encoder
        x, _ = self.decoder(x)                       # BiLSTM decoder -> (B, T, D_MODEL)
        x = x.mean(dim=1)                             # temporal average pool -> (B, D_MODEL)
        return self.out(self.head(x))


# ---------------------------------------------------------------------------
# Train / evaluate (mirrors model.py's train(), swapping in the encoder-decoder model)
# ---------------------------------------------------------------------------


def train(data_dir: Path, coarse: bool = False, num_layers: int = N_ATTN_BLOCKS) -> dict:
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    saved_n, taxonomy = resolve_taxonomy(data_dir, coarse)
    num_classes = len(taxonomy)
    Xtr, ytr = load_split(data_dir, "train")
    Xva, yva = load_split(data_dir, "val")
    Xte, yte = load_split(data_dir, "test")

    assert_label_consistency(saved_n, ytr, yva, yte)
    if coarse:
        def coarsen(y):
            return np.array([labels_mod.FINE_TO_COARSE_IDX[int(v)] for v in y], dtype=np.int64)
        ytr, yva, yte = coarsen(ytr), coarsen(yva), coarsen(yte)
        assert_label_consistency(num_classes, ytr, yva, yte)
    in_features = Xtr.shape[-1]
    label_space_name = "COARSE action-level" if coarse else ("FS_JUMP3D" if taxonomy is labels_mod.FS_JUMP3D_TAXONOMY else "FINE")
    print(f"[Transformer-BiLSTM] label space: {label_space_name} | {num_classes} classes")

    mu = Xtr.mean(axis=(0, 1), keepdims=True)
    sd = Xtr.std(axis=(0, 1), keepdims=True) + 1e-6
    Xtr, Xva, Xte = (Xtr - mu) / sd, (Xva - mu) / sd, (Xte - mu) / sd

    print(f"[Transformer-BiLSTM] device={DEVICE} | num_classes={num_classes} | in_features={in_features}")
    print(f"[Transformer-BiLSTM] shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}")
    print(f"[Transformer-BiLSTM] train classes present: {sorted(set(ytr.tolist()))}")

    present = np.unique(ytr)
    cw = np.clip(compute_class_weight(class_weight="balanced", classes=present, y=ytr),
                 None, MAX_CLASS_WEIGHT)
    weight = torch.ones(num_classes)
    for c, w in zip(present, cw):
        weight[int(c)] = float(w)
    criterion = nn.CrossEntropyLoss(weight=weight.to(DEVICE))

    model = SkatingTransformerBiLSTMClassifier(in_features, num_classes, num_layers=num_layers).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[Transformer-BiLSTM] model params: {n_params/1e6:.2f}M | encoder layers: {num_layers}")
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

    tr_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr)),
        batch_size=BATCH_SIZE, shuffle=True, drop_last=False,
    )
    va_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xva), torch.from_numpy(yva)),
        batch_size=BATCH_SIZE, shuffle=False,
    )

    best_val_f1, best_state, since_improved = -1.0, None, 0
    for epoch in range(1, EPOCHS + 1):
        tr_loss, tr_acc = run_epoch(model, tr_loader, criterion, optimizer)
        va_loss, va_acc = run_epoch(model, va_loader, criterion)
        va_f1 = f1_score(yva, predict(model, Xva), labels=list(range(num_classes)),
                         average="macro", zero_division=0)
        if va_f1 > best_val_f1:
            best_val_f1, since_improved = va_f1, 0
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        else:
            since_improved += 1
        if epoch % 5 == 0 or epoch == 1:
            print(f"[Transformer-BiLSTM] epoch {epoch:3d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} "
                  f"| val loss {va_loss:.3f} acc {va_acc:.3f} f1(macro) {va_f1:.3f}")
        if since_improved >= EARLY_STOP_PATIENCE:
            print(f"[Transformer-BiLSTM] early stop at epoch {epoch} (no val-F1 improvement for {EARLY_STOP_PATIENCE})")
            break

    if best_state is not None:
        model.load_state_dict(best_state)
    print(f"\n[Transformer-BiLSTM] restored best model (val macro-F1 = {best_val_f1:.3f})")

    metrics = report_metrics(yte, predict(model, Xte), taxonomy)
    layer_suffix = "" if num_layers == N_ATTN_BLOCKS else f"_{num_layers}L"
    ckpt_name = f"model_transformer_bilstm{'_coarse' if coarse else ''}{layer_suffix}.pt"
    torch.save({"state_dict": model.state_dict(), "num_classes": num_classes,
                "in_features": in_features, "taxonomy": taxonomy, "coarse": coarse,
                "num_layers": num_layers,
                "feature_mean": mu, "feature_std": sd},
               data_dir / ckpt_name)
    print(f"\n[Transformer-BiLSTM] saved model -> {data_dir / ckpt_name}")
    return metrics


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, required=True)
    parser.add_argument("--coarse", action="store_true",
                        help="collapse jump rotations into action-level classes (11 instead of 28)")
    parser.add_argument("--layers", type=int, default=N_ATTN_BLOCKS,
                        help=f"number of Transformer encoder layers (default {N_ATTN_BLOCKS})")
    args = parser.parse_args()
    if not (args.data_dir / "train_features.pkl").exists():
        raise SystemExit(f"No processed data at {args.data_dir}. Run the pipeline first.")
    train(args.data_dir, coarse=args.coarse, num_layers=args.layers)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())