15_07_expt/scripts/model_bilstm.py

"""CNN-BiLSTM baseline for figure-skating action classification.

Same Conv1D residual backbone and dense head as model.py, but the temporal-modeling stage
is a bidirectional LSTM stack instead of self-attention. Used as an apples-to-apples
comparison against model.py (no positional encoding) and model_transformers.py (self-attention
+ sinusoidal positional encoding) under identical training settings (LR, batch size, class
weighting, early stopping, etc.) -- only the temporal block differs.

Usage:
    python model_bilstm.py --data-dir /path/to/processed
    python model_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 (
    ConvBlock,
    DenseBlock,
    load_split,
    resolve_num_classes,
    assert_label_consistency,
    report_metrics,
    run_epoch,
    predict,
)

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

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"

LSTM_LAYERS = 2
LSTM_DROPOUT = 0.2


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


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

    Same conv backbone + dense head as SkatingActionClassifier (model.py); the attention
    stack is replaced by a 2-layer bidirectional LSTM (hidden=192 each direction -> 384,
    matching the attention stack's channel width so the head is unchanged).
    """

    def __init__(self, in_features: int, num_classes: int):
        super().__init__()
        self.stem = nn.Conv1d(in_features, 128, 3, padding="same")
        self.stem_bn = nn.BatchNorm1d(128)
        self.cb1 = ConvBlock(128, 192, 3)
        self.cb2 = ConvBlock(192, 256, 3)
        self.cb3 = ConvBlock(256, 384, 5)
        self.lstm = nn.LSTM(
            384, 192, num_layers=LSTM_LAYERS, batch_first=True,
            bidirectional=True, dropout=LSTM_DROPOUT,
        )
        self.head = nn.Sequential(
            DenseBlock(384, 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 = x.transpose(1, 2)                      # (B, F, T)
        x = torch.relu(self.stem(x))
        x = self.stem_bn(x)
        x = self.cb3(self.cb2(self.cb1(x)))        # (B, 384, T)
        x = x.transpose(1, 2)                      # (B, T, 384)
        x, _ = self.lstm(x)                         # (B, T, 384)
        x = x.mean(dim=1)                           # temporal average pool -> (B, 384)
        return self.out(self.head(x))


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


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

    fine_n = resolve_num_classes(data_dir)
    taxonomy = labels_mod.COARSE_TAXONOMY if coarse else labels_mod.TAXONOMY
    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(fine_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]
    print(f"[BiLSTM] label space: {'COARSE action-level' if coarse else 'FINE'} | {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"[BiLSTM] device={DEVICE} | num_classes={num_classes} | in_features={in_features}")
    print(f"[BiLSTM] shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}")
    print(f"[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 = SkatingBiLSTMClassifier(in_features, num_classes).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[BiLSTM] model params: {n_params/1e6:.2f}M")
    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"[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"[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[BiLSTM] restored best model (val macro-F1 = {best_val_f1:.3f})")

    metrics = report_metrics(yte, predict(model, Xte), taxonomy)
    ckpt_name = "model_bilstm_coarse.pt" if coarse else "model_bilstm.pt"
    torch.save({"state_dict": model.state_dict(), "num_classes": num_classes,
                "in_features": in_features, "taxonomy": taxonomy, "coarse": coarse,
                "feature_mean": mu, "feature_std": sd},
               data_dir / ckpt_name)
    print(f"\n[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)")
    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)
    return 0


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