Upload folder using huggingface_hub

README.md

@@ -0,0 +1,69 @@
+# ViT-TRM on Something-Something V2
+
+Extends the [ViT-TRM architecture](https://hf.co/adelabdalla221/vit-trm-hmdb51) from HMDB51 (51 classes) to **Something-Something V2** (174 fine-grained hand-object interaction classes).
+
+## Architecture
+
+```
+Video Frames → ViT (per-frame) → Mean Pool → Positional Encoding
+    → TRM Reasoning (H=2 cycles, L=2 shared layers) → Mean Pool → Classifier (174 classes)
+```
+
+- **Backbone**: `vit_tiny_patch16_224` (ImageNet pretrained)
+- **TRM**: 2 cycles × 2 shared transformer layers, 4 heads (~6M params)
+- **Dataset**: SSv2 — 174 template actions, ~220K videos of hand-object interactions
+
+## Setup
+
+```bash
+pip install torch torchvision pytorch-lightning timm torchmetrics decord
+```
+
+## Getting the Data
+
+**Option A: Local download** from [20BN](https://developer.qualcomm.com/software/ai-datasets/something-something):
+```
+ssv2/
+  videos/        # .webm files (1.webm, 2.webm, ...)
+  labels/
+    train.json
+    validation.json
+    labels.json
+```
+
+**Option B: HF Hub** (requires access): `HuggingFaceM4/something-something-v2`
+
+## Training
+
+```bash
+# From scratch
+python train_ssv2.py --data_dir /path/to/ssv2
+
+# Transfer learning from HMDB51 checkpoint (recommended)
+python train_ssv2.py \
+    --data_dir /path/to/ssv2 \
+    --pretrained_ckpt ../vit-trm-hmdb51/vit-trm-epoch=29-val_acc=0.7113.ckpt
+
+# From HF Hub
+python train_ssv2.py --from_hub
+```
+
+### Key flags
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--pretrained_ckpt` | None | Transfer backbone+TRM from HMDB51 |
+| `--trm_H_cycles` | 2 | Number of recursive reasoning cycles |
+| `--frame_stride` | 2 | Temporal stride (SSv2 videos are short) |
+| `--num_frames` | 16 | Frames sampled per clip |
+| `--batch_size` | 8 | Training batch size |
+| `--max_epochs` | 30 | Training epochs |
+| `--precision` | 16-mixed | Mixed precision training |
+
+## Why SSv2?
+
+Unlike HMDB51 which can be solved partly by scene/object appearance, SSv2 requires **temporal reasoning** — understanding the motion and interaction pattern. This makes it a better test of the TRM recursive reasoning approach:
+
+- "Pushing something from left to right" vs "Pushing something from right to left" differ only in motion direction
+- 174 fine-grained template actions, ~220K training videos
+- Standard benchmark for temporal modeling in video understanding

ssv2_datamodule.py

@@ -0,0 +1,362 @@
+#!/usr/bin/env python3
+"""
+Something-Something V2 DataModule for PyTorch Lightning.
+
+Loads SSv2 from the Hugging Face Hub or from a local directory of webm files.
+Each sample is a short video (~2-6 s) of a hand performing one of 174 template actions
+(e.g. "Pushing [something] from left to right").
+
+Usage:
+    dm = SSv2DataModule(data_dir="/path/to/ssv2", batch_size=8)
+    dm.setup()
+    for batch in dm.train_dataloader():
+        ...
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional, Callable, List, Dict, Tuple
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+import pytorch_lightning as pl
+import torchvision.transforms as T
+
+try:
+    import decord
+    decord.bridge.set_bridge("torch")
+    HAS_DECORD = True
+except ImportError:
+    HAS_DECORD = False
+
+try:
+    from datasets import load_dataset as hf_load_dataset
+    HAS_HF_DATASETS = True
+except ImportError:
+    HAS_HF_DATASETS = False
+
+
+# ---------------------------------------------------------------------------
+# Video sampling helpers
+# ---------------------------------------------------------------------------
+
+def sample_frames_uniform(total_frames: int, num_frames: int) -> List[int]:
+    """Uniformly sample `num_frames` indices from [0, total_frames)."""
+    if total_frames <= num_frames:
+        indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames)
+        return indices
+    stride = total_frames / num_frames
+    return [int(i * stride) for i in range(num_frames)]
+
+
+def sample_frames_stride(total_frames: int, num_frames: int, stride: int) -> List[int]:
+    """Sample `num_frames` with fixed stride, centered in the video."""
+    needed = (num_frames - 1) * stride + 1
+    if needed > total_frames:
+        return sample_frames_uniform(total_frames, num_frames)
+    start = (total_frames - needed) // 2
+    return [start + i * stride for i in range(num_frames)]
+
+
+# ---------------------------------------------------------------------------
+# Dataset: local directory of webm/mp4 files + label JSON
+# ---------------------------------------------------------------------------
+
+class SSv2LocalDataset(Dataset):
+    """
+    Loads SSv2 from a local directory.
+
+    Expected layout:
+        data_dir/
+            videos/           # or 20bn-something-something-v2/
+                1.webm
+                2.webm
+                ...
+            labels/
+                train.json    # [{"id": "1", "template": "...", "label": "..."}, ...]
+                validation.json
+                test.json     # (no labels)
+                labels.json   # {"0": "Approaching [something] with your camera", ...}
+    """
+
+    def __init__(
+        self,
+        data_dir: str,
+        split: str = "train",
+        num_frames: int = 16,
+        frame_stride: int = 2,
+        transform: Optional[Callable] = None,
+        num_clips: int = 1,
+    ):
+        super().__init__()
+        self.data_dir = Path(data_dir)
+        self.split = split
+        self.num_frames = num_frames
+        self.frame_stride = frame_stride
+        self.transform = transform
+        self.num_clips = num_clips
+
+        if not HAS_DECORD:
+            raise ImportError("decord is required for local video loading. Install: pip install decord")
+
+        # Find video directory
+        vid_dirs = ["videos", "20bn-something-something-v2"]
+        self.video_dir = None
+        for d in vid_dirs:
+            candidate = self.data_dir / d
+            if candidate.exists():
+                self.video_dir = candidate
+                break
+        if self.video_dir is None:
+            self.video_dir = self.data_dir / "videos"
+
+        # Load labels mapping
+        labels_file = self.data_dir / "labels" / "labels.json"
+        if labels_file.exists():
+            with open(labels_file) as f:
+                idx_to_label = json.load(f)
+            self.label_to_idx = {v: int(k) for k, v in idx_to_label.items()}
+        else:
+            self.label_to_idx = {}
+
+        # Load split annotations
+        split_file = self.data_dir / "labels" / f"{split}.json"
+        if not split_file.exists():
+            # Try alternate naming
+            alt = self.data_dir / "labels" / f"something-something-v2-{split}.json"
+            if alt.exists():
+                split_file = alt
+            else:
+                raise FileNotFoundError(f"Cannot find annotation file for split '{split}' in {self.data_dir / 'labels'}")
+
+        with open(split_file) as f:
+            self.annotations = json.load(f)
+
+        # Build label_to_idx from annotations if not loaded from labels.json
+        if not self.label_to_idx:
+            all_labels = sorted(set(a.get("template", a.get("label", "")) for a in self.annotations if "template" in a or "label" in a))
+            self.label_to_idx = {lbl: i for i, lbl in enumerate(all_labels)}
+
+        self.num_classes = len(self.label_to_idx)
+        print(f"SSv2 [{split}]: {len(self.annotations)} videos, {self.num_classes} classes")
+
+    def __len__(self):
+        return len(self.annotations) * self.num_clips
+
+    def __getitem__(self, idx):
+        clip_idx = idx % self.num_clips
+        video_idx = idx // self.num_clips
+        ann = self.annotations[video_idx]
+
+        video_id = str(ann["id"])
+        label_str = ann.get("template", ann.get("label", None))
+        label = self.label_to_idx.get(label_str, -1) if label_str else -1
+
+        # Find video file
+        video_path = None
+        for ext in [".webm", ".mp4"]:
+            candidate = self.video_dir / f"{video_id}{ext}"
+            if candidate.exists():
+                video_path = str(candidate)
+                break
+        if video_path is None:
+            raise FileNotFoundError(f"Video not found: {video_id} in {self.video_dir}")
+
+        # Decode frames
+        vr = decord.VideoReader(video_path)
+        total = len(vr)
+        indices = sample_frames_stride(total, self.num_frames, self.frame_stride)
+        frames = vr.get_batch(indices)  # (T, H, W, C) as torch tensor
+
+        # Convert to (T, C, H, W) float [0,1]
+        frames = frames.permute(0, 3, 1, 2).float() / 255.0
+
+        if self.transform is not None:
+            frames = torch.stack([self.transform(f) for f in frames])
+
+        return {"video": frames, "label": label, "video_id": video_id}
+
+
+# ---------------------------------------------------------------------------
+# Dataset: Hugging Face Hub streaming
+# ---------------------------------------------------------------------------
+
+class SSv2HFDataset(Dataset):
+    """
+    Loads SSv2 from the Hugging Face Hub using the `datasets` library.
+    Tries known Hub IDs: 'HuggingFaceM4/something-something-v2' or 'lmms-lab/SSv2'.
+    Falls back to manual download instructions if gated.
+    """
+
+    def __init__(
+        self,
+        split: str = "train",
+        num_frames: int = 16,
+        frame_stride: int = 2,
+        transform: Optional[Callable] = None,
+        num_clips: int = 1,
+        hf_dataset_id: str = "HuggingFaceM4/something-something-v2",
+    ):
+        super().__init__()
+        if not HAS_HF_DATASETS:
+            raise ImportError("Install: pip install datasets")
+
+        self.num_frames = num_frames
+        self.frame_stride = frame_stride
+        self.transform = transform
+        self.num_clips = num_clips
+
+        print(f"Loading SSv2 from Hub: {hf_dataset_id} (split={split}) ...")
+        self.ds = hf_load_dataset(hf_dataset_id, split=split)
+
+        # Infer label column and build mapping
+        if "label" in self.ds.features:
+            feat = self.ds.features["label"]
+            if hasattr(feat, "names"):
+                self.num_classes = len(feat.names)
+            else:
+                self.num_classes = 174
+        else:
+            self.num_classes = 174
+
+        print(f"SSv2 HF [{split}]: {len(self.ds)} samples, {self.num_classes} classes")
+
+    def __len__(self):
+        return len(self.ds) * self.num_clips
+
+    def __getitem__(self, idx):
+        video_idx = idx // self.num_clips
+        sample = self.ds[video_idx]
+
+        label = sample.get("label", -1)
+        video_id = str(sample.get("video_id", sample.get("id", video_idx)))
+
+        # The HF dataset typically stores video as bytes or decoded frames
+        video_data = sample.get("video", None)
+        if video_data is None:
+            raise ValueError("No 'video' column in HF dataset")
+
+        # If video_data is a dict with 'path'/'bytes', decode with decord
+        if isinstance(video_data, dict):
+            import io
+            video_bytes = video_data.get("bytes", None)
+            if video_bytes:
+                vr = decord.VideoReader(io.BytesIO(video_bytes))
+                total = len(vr)
+                indices = sample_frames_stride(total, self.num_frames, self.frame_stride)
+                frames = vr.get_batch(indices).permute(0, 3, 1, 2).float() / 255.0
+            else:
+                raise ValueError("Cannot decode video from HF dataset sample")
+        elif isinstance(video_data, torch.Tensor):
+            frames = video_data
+            if frames.ndim == 4 and frames.shape[-1] in (1, 3):
+                frames = frames.permute(0, 3, 1, 2).float()
+                if frames.max() > 1.0:
+                    frames = frames / 255.0
+            total = frames.shape[0]
+            indices = sample_frames_stride(total, self.num_frames, self.frame_stride)
+            frames = frames[indices]
+        else:
+            raise ValueError(f"Unexpected video format: {type(video_data)}")
+
+        if self.transform is not None:
+            frames = torch.stack([self.transform(f) for f in frames])
+
+        return {"video": frames, "label": label, "video_id": video_id}
+
+
+# ---------------------------------------------------------------------------
+# Lightning DataModule
+# ---------------------------------------------------------------------------
+
+def build_train_transform(img_size: int = 224):
+    return T.Compose([
+        T.Resize((img_size, img_size)),
+        T.RandomHorizontalFlip(),
+        T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+
+def build_val_transform(img_size: int = 224):
+    return T.Compose([
+        T.Resize((img_size, img_size)),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+
+class SSv2DataModule(pl.LightningDataModule):
+    """
+    SSv2 DataModule supporting both local files and HF Hub.
+
+    Args:
+        data_dir: Path to local SSv2 data. If None, loads from HF Hub.
+        hf_dataset_id: HF Hub dataset ID (used when data_dir is None).
+        num_frames: Frames to sample per clip.
+        frame_stride: Temporal stride between sampled frames.
+        img_size: Spatial resize target.
+        batch_size: Training batch size.
+        num_workers: DataLoader workers.
+        num_clips_val: Number of clips per video at val/test time.
+    """
+
+    def __init__(
+        self,
+        data_dir: Optional[str] = None,
+        hf_dataset_id: str = "HuggingFaceM4/something-something-v2",
+        num_frames: int = 16,
+        frame_stride: int = 2,
+        img_size: int = 224,
+        batch_size: int = 8,
+        num_workers: int = 4,
+        num_clips_val: int = 4,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.data_dir = data_dir
+        self.hf_dataset_id = hf_dataset_id
+        self.num_frames = num_frames
+        self.frame_stride = frame_stride
+        self.img_size = img_size
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.num_clips_val = num_clips_val
+
+    def setup(self, stage=None):
+        train_tf = build_train_transform(self.img_size)
+        val_tf = build_val_transform(self.img_size)
+
+        if self.data_dir is not None:
+            self.train_ds = SSv2LocalDataset(
+                self.data_dir, "train", self.num_frames, self.frame_stride, train_tf, num_clips=1,
+            )
+            self.val_ds = SSv2LocalDataset(
+                self.data_dir, "validation", self.num_frames, self.frame_stride, val_tf, num_clips=self.num_clips_val,
+            )
+            self.num_classes = self.train_ds.num_classes
+        else:
+            self.train_ds = SSv2HFDataset(
+                "train", self.num_frames, self.frame_stride, train_tf, num_clips=1,
+                hf_dataset_id=self.hf_dataset_id,
+            )
+            self.val_ds = SSv2HFDataset(
+                "validation", self.num_frames, self.frame_stride, val_tf, num_clips=self.num_clips_val,
+                hf_dataset_id=self.hf_dataset_id,
+            )
+            self.num_classes = self.train_ds.num_classes
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_ds, batch_size=self.batch_size, shuffle=True,
+            num_workers=self.num_workers, pin_memory=True, drop_last=True,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_ds, batch_size=self.batch_size, shuffle=False,
+            num_workers=self.num_workers, pin_memory=True,
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()

train_ssv2.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""
+Train ViT-TRM on Something-Something V2.
+
+Examples:
+    # From scratch on local SSv2 data:
+    python train_ssv2.py --data_dir /path/to/ssv2
+
+    # Transfer from HMDB51 pretrained checkpoint:
+    python train_ssv2.py --data_dir /path/to/ssv2 --pretrained_ckpt ../vit-trm-hmdb51/vit-trm-epoch=29-val_acc=0.7113.ckpt
+
+    # From HF Hub (if you have access):
+    python train_ssv2.py --from_hub
+
+    # Quick smoke test (2 epochs, 1 batch):
+    python train_ssv2.py --data_dir /path/to/ssv2 --fast_dev_run
+"""
+
+import argparse
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+
+from vit_trm_video import ViTTRMVideo
+from ssv2_datamodule import SSv2DataModule
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train ViT-TRM on SSv2")
+
+    # Data
+    parser.add_argument("--data_dir", type=str, default=None, help="Local SSv2 data directory")
+    parser.add_argument("--from_hub", action="store_true", help="Load SSv2 from HF Hub")
+    parser.add_argument("--hf_dataset_id", type=str, default="HuggingFaceM4/something-something-v2")
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--frame_stride", type=int, default=2, help="SSv2 videos are short, use stride=2")
+    parser.add_argument("--img_size", type=int, default=224)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--num_clips_val", type=int, default=4)
+
+    # Model
+    parser.add_argument("--vit_name", type=str, default="vit_tiny_patch16_224")
+    parser.add_argument("--vit_pretrained", action="store_true", default=True)
+    parser.add_argument("--vit_freeze", action="store_true", default=False)
+    parser.add_argument("--trm_H_cycles", type=int, default=2)
+    parser.add_argument("--trm_L_layers", type=int, default=2)
+    parser.add_argument("--trm_num_heads", type=int, default=4)
+    parser.add_argument("--num_classes", type=int, default=174)
+    parser.add_argument("--pretrained_ckpt", type=str, default=None,
+                        help="Path to HMDB51 checkpoint to transfer backbone+TRM weights from")
+
+    # Training
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--weight_decay", type=float, default=0.05)
+    parser.add_argument("--warmup_epochs", type=int, default=5)
+    parser.add_argument("--max_epochs", type=int, default=30)
+    parser.add_argument("--label_smoothing", type=float, default=0.1)
+    parser.add_argument("--iterative_refinement", action="store_true", default=False)
+
+    # Trainer
+    parser.add_argument("--accelerator", type=str, default="auto")
+    parser.add_argument("--devices", type=int, default=1)
+    parser.add_argument("--precision", type=str, default="16-mixed")
+    parser.add_argument("--fast_dev_run", action="store_true", default=False)
+    parser.add_argument("--seed", type=int, default=42)
+
+    args = parser.parse_args()
+    pl.seed_everything(args.seed)
+
+    # Data
+    data_dir = args.data_dir if not args.from_hub else None
+    dm = SSv2DataModule(
+        data_dir=data_dir,
+        hf_dataset_id=args.hf_dataset_id,
+        num_frames=args.num_frames,
+        frame_stride=args.frame_stride,
+        img_size=args.img_size,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        num_clips_val=args.num_clips_val,
+    )
+
+    # Model
+    model = ViTTRMVideo(
+        img_size=args.img_size,
+        vit_name=args.vit_name,
+        vit_pretrained=args.vit_pretrained,
+        vit_freeze=args.vit_freeze,
+        trm_H_cycles=args.trm_H_cycles,
+        trm_L_layers=args.trm_L_layers,
+        trm_num_heads=args.trm_num_heads,
+        num_classes=args.num_classes,
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        warmup_epochs=args.warmup_epochs,
+        max_epochs=args.max_epochs,
+        label_smoothing=args.label_smoothing,
+        iterative_refinement=args.iterative_refinement,
+        pretrained_ckpt=args.pretrained_ckpt,
+    )
+
+    # Callbacks
+    ckpt_callback = ModelCheckpoint(
+        dirpath="checkpoints",
+        filename="vit-trm-ssv2-{epoch:02d}-{val_acc:.4f}",
+        monitor="val_acc",
+        mode="max",
+        save_top_k=3,
+    )
+    lr_monitor = LearningRateMonitor(logging_interval="epoch")
+
+    # Trainer
+    trainer = pl.Trainer(
+        accelerator=args.accelerator,
+        devices=args.devices,
+        precision=args.precision,
+        max_epochs=args.max_epochs,
+        callbacks=[ckpt_callback, lr_monitor],
+        fast_dev_run=args.fast_dev_run,
+        log_every_n_steps=50,
+    )
+
+    trainer.fit(model, dm)
+
+    # Test with best checkpoint
+    if not args.fast_dev_run:
+        trainer.test(model, dm, ckpt_path="best")
+
+
+if __name__ == "__main__":
+    main()

vit_trm_video.py

@@ -0,0 +1,348 @@
+#!/usr/bin/env python3
+"""
+ViT + TRM Video Classifier — dataset-agnostic version.
+
+Architecture:
+- ViT per-frame feature extraction
+- TRM reasoning cycles (shared-weight transformer layers)
+- Temporal pooling
+- Classifier
+
+Supports video-level evaluation by aggregating multi-clip predictions.
+"""
+
+from typing import Optional, Dict
+import math
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+import timm
+
+
+def build_sinusoidal_positional_encoding(seq_len: int, dim: int, device: torch.device) -> torch.Tensor:
+    position = torch.arange(seq_len, device=device).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, dim, 2, device=device) * (-torch.log(torch.tensor(10000.0)) / dim))
+    pe = torch.zeros(seq_len, dim, device=device)
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    return pe.unsqueeze(0)
+
+
+class ReasoningCycle(nn.Module):
+    """
+    Single reasoning cycle (TRM's H-cycle).
+    Applies L shared transformer layers to refine representations.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, num_layers: int, dropout: float = 0.1):
+        super().__init__()
+        self.num_layers = num_layers
+        self.shared_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            nhead=num_heads,
+            dim_feedforward=hidden_size * 4,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for _ in range(self.num_layers):
+            x = self.shared_layer(x)
+        return self.norm(x)
+
+
+class ViTTRMVideo(pl.LightningModule):
+    """
+    ViT + TRM for video classification.
+
+    Architecture:
+    1. ViT per-frame feature extraction
+    2. TRM recursive reasoning over temporal tokens
+    3. Mean-pool + Classifier
+    """
+
+    def __init__(
+        self,
+        # Frame encoder (ViT) config
+        img_size: int = 224,
+        vit_name: str = "vit_tiny_patch16_224",
+        vit_pretrained: bool = True,
+        vit_freeze: bool = False,
+        # TRM config
+        trm_H_cycles: int = 2,
+        trm_L_layers: int = 2,
+        trm_hidden_size: Optional[int] = None,
+        trm_num_heads: int = 4,
+        # Task config
+        num_classes: int = 174,
+        # Training config
+        lr: float = 3e-4,
+        weight_decay: float = 0.05,
+        warmup_epochs: int = 5,
+        max_epochs: int = 50,
+        label_smoothing: float = 0.1,
+        # Iterative refinement
+        iterative_refinement: bool = False,
+        num_refinement_steps: int = None,
+        # Transfer learning — path to a pretrained checkpoint to load backbone + TRM from
+        pretrained_ckpt: Optional[str] = None,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.warmup_epochs = warmup_epochs
+        self.max_epochs = max_epochs
+        self.num_classes = num_classes
+        self.label_smoothing = label_smoothing
+        self.trm_H_cycles = trm_H_cycles
+        self.iterative_refinement = iterative_refinement
+
+        if num_refinement_steps is None:
+            self.num_refinement_steps = trm_H_cycles
+        else:
+            self.num_refinement_steps = num_refinement_steps
+
+        if iterative_refinement:
+            self.automatic_optimization = False
+
+        # ViT backbone
+        self.vit = timm.create_model(
+            vit_name,
+            pretrained=vit_pretrained,
+            num_classes=0,
+            img_size=img_size,
+            dynamic_img_size=True,
+        )
+        if hasattr(self.vit, "reset_classifier"):
+            self.vit.reset_classifier(0, global_pool="")
+
+        self.vit_freeze = vit_freeze
+        if vit_freeze:
+            for p in self.vit.parameters():
+                p.requires_grad = False
+            self.vit.eval()
+
+        vit_embed_dim = getattr(self.vit, "num_features", None) or getattr(self.vit, "embed_dim", None)
+        if vit_embed_dim is None:
+            raise ValueError("Could not infer ViT embedding dimension from timm model.")
+
+        if trm_hidden_size is None:
+            trm_hidden_size = int(vit_embed_dim)
+        self.trm_hidden_size = trm_hidden_size
+
+        # TRM reasoning cycles
+        self.reasoning_cycle = ReasoningCycle(
+            hidden_size=self.trm_hidden_size,
+            num_heads=trm_num_heads,
+            num_layers=trm_L_layers,
+            dropout=0.1,
+        )
+
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.LayerNorm(self.trm_hidden_size),
+            nn.Linear(self.trm_hidden_size, num_classes),
+        )
+
+        # Metrics
+        import torchmetrics
+        self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
+        self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
+
+        self.validation_outputs = []
+
+        # Optionally load pretrained weights (e.g. from HMDB51 checkpoint)
+        if pretrained_ckpt is not None:
+            self._load_pretrained(pretrained_ckpt)
+
+    def _load_pretrained(self, ckpt_path: str):
+        """Load backbone + TRM weights from a prior checkpoint, skip classifier."""
+        ckpt = torch.load(ckpt_path, map_location="cpu")
+        state_dict = ckpt.get("state_dict", ckpt)
+        # Filter out classifier weights (different num_classes)
+        filtered = {k: v for k, v in state_dict.items() if not k.startswith("classifier.")}
+        missing, unexpected = self.load_state_dict(filtered, strict=False)
+        print(f"Loaded pretrained weights from {ckpt_path}")
+        print(f"  Missing keys (expected — new classifier): {missing}")
+        if unexpected:
+            print(f"  Unexpected keys (ignored): {unexpected}")
+
+    def forward(self, video, num_cycles=None):
+        if num_cycles is None:
+            num_cycles = self.trm_H_cycles
+
+        B, T, C, H, W = video.shape
+        frames_bt = video.view(B * T, C, H, W)
+
+        tokens = self.vit.forward_features(frames_bt)
+        frame_features = tokens.mean(dim=1)
+        features = frame_features.view(B, T, -1)
+        pos = build_sinusoidal_positional_encoding(T, features.size(-1), features.device)
+        features = features + pos
+
+        if num_cycles > 0:
+            for _ in range(num_cycles):
+                features = self.reasoning_cycle(features)
+
+        pooled = features.mean(dim=1)
+        logits = self.classifier(pooled)
+        return logits
+
+    def _unpack_batch(self, batch: Dict[str, torch.Tensor]):
+        if isinstance(batch, tuple):
+            return batch[0], batch[1], None
+        video_ids = batch.get("video_id", None)
+        return batch["video"], batch["label"], video_ids
+
+    def training_step(self, batch, batch_idx):
+        videos, labels, _ = self._unpack_batch(batch)
+
+        if self.iterative_refinement:
+            opt = self.optimizers()
+            opt.zero_grad()
+            total_loss = 0.0
+            for step in range(1, self.num_refinement_steps + 1):
+                logits = self(videos, num_cycles=step)
+                loss = nn.functional.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
+                total_loss += loss / self.num_refinement_steps
+            self.manual_backward(total_loss)
+            opt.step()
+            with torch.no_grad():
+                final_logits = self(videos)
+                final_preds = torch.argmax(final_logits, dim=1)
+            sch = self.lr_schedulers()
+            if sch is not None:
+                sch.step()
+            acc = self.train_acc(final_preds, labels)
+            self.log("train_loss", total_loss, on_step=True, on_epoch=True, prog_bar=True)
+            self.log("train_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+            return total_loss
+        else:
+            logits = self(videos)
+            loss = nn.functional.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
+            preds = torch.argmax(logits, dim=1)
+            acc = self.train_acc(preds, labels)
+            self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
+            self.log("train_acc", acc, on_step=False, on_epoch=True, prog_bar=True)
+            return loss
+
+    def validation_step(self, batch, batch_idx):
+        videos, labels, video_ids = self._unpack_batch(batch)
+        logits = self(videos)
+        loss = nn.functional.cross_entropy(logits, labels)
+        preds = torch.argmax(logits, dim=1)
+        acc = self.val_acc(preds, labels)
+        self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("val_acc_clip", acc, on_step=False, on_epoch=True, prog_bar=True)
+        if video_ids is not None:
+            self.validation_outputs.append({
+                "video_ids": video_ids,
+                "logits": logits.detach().cpu(),
+                "labels": labels.detach().cpu(),
+                "preds": preds.detach().cpu(),
+            })
+        return loss
+
+    def on_validation_epoch_end(self):
+        if not self.validation_outputs:
+            return
+        from collections import defaultdict
+        video_logits = defaultdict(list)
+        video_labels = {}
+        for output in self.validation_outputs:
+            for i, vid in enumerate(output["video_ids"]):
+                video_logits[vid].append(output["logits"][i])
+                video_labels[vid] = output["labels"][i].item()
+        video_preds, video_true = [], []
+        for vid in sorted(video_logits.keys()):
+            avg = torch.stack(video_logits[vid]).mean(dim=0)
+            video_preds.append(torch.argmax(avg).item())
+            video_true.append(video_labels[vid])
+        video_acc = (torch.tensor(video_preds) == torch.tensor(video_true)).float().mean()
+        self.log("val_acc_video", video_acc, on_epoch=True, prog_bar=True)
+        self.log("val_acc", video_acc, on_epoch=True, prog_bar=True)
+        num_videos = len(video_logits)
+        num_clips = sum(len(v) for v in video_logits.values())
+        print(f"\n  Video-level val: {num_videos} videos, {num_clips} clips, acc={video_acc:.4f}")
+        self.validation_outputs.clear()
+
+    def test_step(self, batch, batch_idx):
+        videos, labels, video_ids = self._unpack_batch(batch)
+        logits = self(videos)
+        loss = nn.functional.cross_entropy(logits, labels)
+        preds = torch.argmax(logits, dim=1)
+        acc = self.val_acc(preds, labels)
+        self.log("test_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("test_acc_clip", acc, on_step=False, on_epoch=True, prog_bar=True)
+        if video_ids is not None:
+            self.validation_outputs.append({
+                "video_ids": video_ids,
+                "logits": logits.detach().cpu(),
+                "labels": labels.detach().cpu(),
+                "preds": preds.detach().cpu(),
+            })
+        return loss
+
+    def on_test_epoch_end(self):
+        if not self.validation_outputs:
+            return
+        from collections import defaultdict
+        video_logits = defaultdict(list)
+        video_labels = {}
+        for output in self.validation_outputs:
+            for i, vid in enumerate(output["video_ids"]):
+                video_logits[vid].append(output["logits"][i])
+                video_labels[vid] = output["labels"][i].item()
+        video_preds, video_true = [], []
+        for vid in sorted(video_logits.keys()):
+            avg = torch.stack(video_logits[vid]).mean(dim=0)
+            video_preds.append(torch.argmax(avg).item())
+            video_true.append(video_labels[vid])
+        video_acc = (torch.tensor(video_preds) == torch.tensor(video_true)).float().mean()
+        self.log("test_acc_video", video_acc, on_epoch=True, prog_bar=True)
+        self.log("test_acc", video_acc, on_epoch=True, prog_bar=True)
+        print(f"\n  Video-level test: {len(video_logits)} videos, acc={video_acc:.4f}")
+        self.validation_outputs.clear()
+
+    def on_train_epoch_start(self):
+        if self.vit_freeze:
+            self.vit.eval()
+
+    def configure_optimizers(self):
+        decay, no_decay = [], []
+        for n, p in self.named_parameters():
+            if not p.requires_grad:
+                continue
+            if p.ndim < 2 or n.endswith("bias") or "norm" in n.lower() or "bn" in n.lower():
+                no_decay.append(p)
+            else:
+                decay.append(p)
+        vit_param_ids = {id(p) for p in self.vit.parameters()}
+        optimizer = torch.optim.AdamW([
+            {"params": [p for p in decay if id(p) not in vit_param_ids], "lr": self.lr, "weight_decay": self.weight_decay},
+            {"params": [p for p in no_decay if id(p) not in vit_param_ids], "lr": self.lr, "weight_decay": 0.0},
+            {"params": [p for p in decay if id(p) in vit_param_ids], "lr": self.lr * 0.1, "weight_decay": self.weight_decay},
+            {"params": [p for p in no_decay if id(p) in vit_param_ids], "lr": self.lr * 0.1, "weight_decay": 0.0},
+        ])
+
+        def lr_lambda(epoch: int) -> float:
+            if epoch < self.warmup_epochs:
+                return float((epoch + 1) / max(1, self.warmup_epochs))
+            progress = (epoch - self.warmup_epochs) / max(1, (self.max_epochs - self.warmup_epochs))
+            return 0.5 * (1.0 + math.cos(math.pi * min(1.0, max(0.0, progress))))
+
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+        return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}}
+
+
+if __name__ == "__main__":
+    model = ViTTRMVideo(num_classes=174, trm_H_cycles=2)
+    x = torch.randn(2, 16, 3, 224, 224)
+    y = model(x)
+    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Trainable parameters: {num_params:,}")
+    print("Logits:", y.shape)  # (2, 174)