bcgxtberg/vit-trm-ssv2

ViT-TRM on UCF-101

Extends the ViT-TRM architecture from HMDB51 (51 classes) to UCF-101 (101 action classes) via transfer learning.

Architecture

Video Frames → ViT (per-frame) → Mean Pool → Positional Encoding
    → TRM Reasoning (H=2 cycles, L=2 shared layers) → Mean Pool → Classifier (101 classes)

• Backbone: vit_tiny_patch16_224 (ImageNet pretrained)
• TRM: 2 cycles × 2 shared transformer layers, 4 heads (~6M params)
• Dataset: UCF-101 — 101 action classes, 13K+ videos
• Transfer: Backbone + TRM weights initialized from HMDB51 checkpoint (69% accuracy)
• Paper: Less is More: Recursive Reasoning with Tiny Networks

Setup

pip install torch torchvision pytorch-lightning timm torchmetrics datasets

Training

# Transfer from HMDB51 (recommended)
python train_ucf101.py --pretrained_ckpt vit-trm-epoch=29-val_acc=0.7113.ckpt

# From scratch
python train_ucf101.py

# Smoke test
python train_ucf101.py --fast_dev_run --max_videos 50

Usage

import torch
from vit_trm_video import ViTTRMVideo

model = ViTTRMVideo.load_from_checkpoint("checkpoints/best.ckpt", strict=False)
model.eval()

video = torch.randn(1, 16, 3, 224, 224)  # (batch, frames, C, H, W)
with torch.no_grad():
    logits = model(video)
    pred = logits.argmax(dim=-1)