| """LS-ViT model architecture for HMDB51 action recognition. |
| |
| This module defines the LS-ViT (Long-Short ViT) architecture used to train the |
| weights stored in `lsvit_hmdb51_best.pt`. The model wraps a ViT-Base backbone |
| with two motion-aware modules: |
| |
| - SMIFModule: Short-term Motion Injection & Fusion, applied to raw RGB frames. |
| - LMIModule: Long-term Motion Interaction, applied to patch tokens inside |
| every transformer block. |
| |
| Usage: |
| import torch |
| from modeling import ViTConfig, LSViTForAction |
| |
| config = ViTConfig(image_size=224) |
| model = LSViTForAction(config, num_classes=51) |
| ckpt = torch.load("lsvit_hmdb51_best.pt", map_location="cpu", weights_only=False) |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
| |
| # video: (Batch, Time, Channels, Height, Width) in [0, 1] after standard ViT |
| # normalization. Trained with T=12, image_size=224. |
| logits = model(video) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class ViTConfig: |
| image_size: int = 224 |
| patch_size: int = 16 |
| in_chans: int = 3 |
| embed_dim: int = 768 |
| depth: int = 12 |
| num_heads: int = 12 |
| mlp_ratio: float = 4.0 |
| drop_rate: float = 0.1 |
| attn_drop_rate: float = 0.1 |
| drop_path_rate: float = 0.1 |
| qkv_bias: bool = True |
|
|
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, config: ViTConfig): |
| super().__init__() |
| self.image_size = config.image_size |
| self.patch_size = config.patch_size |
| self.num_patches = (config.image_size // config.patch_size) ** 2 |
|
|
| self.proj = nn.Conv2d( |
| config.in_chans, |
| config.embed_dim, |
| kernel_size=config.patch_size, |
| stride=config.patch_size, |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.proj(x) |
| x = x.flatten(2).transpose(1, 2) |
| return x |
|
|
|
|
| class DropPath(nn.Module): |
| """Drop paths per sample when applied in the main path of residual blocks.""" |
|
|
| def __init__(self, drop_prob: float = 0.0): |
| super().__init__() |
| self.drop_prob = drop_prob |
|
|
| def forward(self, x): |
| if self.drop_prob == 0.0 or not self.training: |
| return x |
| keep_prob = 1 - self.drop_prob |
| shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| random_tensor.floor_() |
| return x.div(keep_prob) * random_tensor |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, dim: int, num_heads: int, qkv_bias: bool, attn_drop: float, proj_drop: float): |
| super().__init__() |
| self.num_heads = num_heads |
| head_dim = dim // num_heads |
| self.scale = head_dim ** -0.5 |
|
|
| self.q = nn.Linear(dim, dim, bias=qkv_bias) |
| self.k = nn.Linear(dim, dim, bias=qkv_bias) |
| self.v = nn.Linear(dim, dim, bias=qkv_bias) |
|
|
| self.attn_drop = nn.Dropout(attn_drop) |
| self.proj = nn.Linear(dim, dim) |
| self.proj_drop = nn.Dropout(proj_drop) |
|
|
| def forward(self, x): |
| B, N, C = x.shape |
|
|
| q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
| k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
| v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
|
|
| attn = (q @ k.transpose(-2, -1)) * self.scale |
| attn = attn.softmax(dim=-1) |
| attn = self.attn_drop(attn) |
|
|
| x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| x = self.proj(x) |
| x = self.proj_drop(x) |
| return x |
|
|
|
|
| class SMIFModule(nn.Module): |
| """Short-term Motion Injection & Fusion over raw RGB frames.""" |
|
|
| def __init__(self, channels: int, window_size: int = 5, alpha: float = 0.5, threshold: float = 0.05): |
| super().__init__() |
| assert window_size % 2 == 1, "window_size must be odd" |
| self.channels = channels |
| self.window_size = window_size |
| self.half = window_size // 2 |
| self.threshold = threshold |
|
|
| self.alpha = nn.Parameter(torch.tensor(alpha)) |
| self.conv_fuse = nn.Conv2d(channels * 2, channels, kernel_size=1) |
|
|
| def forward(self, video: torch.Tensor, return_motion_map: bool = False): |
| B, T, C, H, W = video.shape |
|
|
| motion_accum = torch.zeros_like(video) |
| for offset in range(1, self.half + 1): |
| prev_frames = torch.roll(video, shifts=offset, dims=1) |
| next_frames = torch.roll(video, shifts=-offset, dims=1) |
| prev_frames[:, :offset] = video[:, :offset] |
| next_frames[:, -offset:] = video[:, -offset:] |
|
|
| diff_future = next_frames - video |
| diff_past = video - prev_frames |
| motion_accum = motion_accum + diff_future.abs() + diff_past.abs() |
|
|
| motion_map = motion_accum / max(self.half, 1) |
| mask = (motion_map > self.threshold).float() |
| motion_map = motion_map * mask |
|
|
| base_2d = video.reshape(B * T, C, H, W) |
| motion_2d = motion_map.reshape(B * T, C, H, W) |
|
|
| fused = torch.cat([base_2d, motion_2d], dim=1) |
| fused = self.conv_fuse(fused) |
|
|
| out = base_2d + self.alpha.tanh() * fused |
| out = out.clamp(min=-1.0, max=1.0) |
| out = out.view(B, T, C, H, W) |
|
|
| if return_motion_map: |
| return out, motion_map |
| return out |
|
|
|
|
| class LMIModule(nn.Module): |
| """Long-term Motion Interaction operating on token differences across time.""" |
|
|
| def __init__(self, dim: int, reduction: int = 4, delta: float = 0.1): |
| super().__init__() |
| reduced_dim = max(1, dim // reduction) |
| self.reduce = nn.Linear(dim, reduced_dim) |
| self.expand = nn.Linear(reduced_dim, dim) |
| self.temporal_mlp = nn.Sequential( |
| nn.LayerNorm(reduced_dim), |
| nn.Linear(reduced_dim, reduced_dim), |
| nn.GELU(), |
| nn.Linear(reduced_dim, reduced_dim), |
| ) |
| self.delta = nn.Parameter(torch.tensor(delta)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, T, N, C = x.shape |
| reduced = self.reduce(x) |
|
|
| if T > 1: |
| diff_f = reduced[:, 1:] - reduced[:, :-1] |
| diff_f = torch.cat([diff_f, diff_f[:, -1:]], dim=1) |
| diff_b = reduced[:, :-1] - reduced[:, 1:] |
| diff_b = torch.cat([diff_b[:, :1], diff_b], dim=1) |
| else: |
| diff_f = torch.zeros_like(reduced) |
| diff_b = torch.zeros_like(reduced) |
|
|
| motion = (diff_f.abs() + diff_b.abs()).mean(dim=2) |
| motion = self.temporal_mlp(motion) |
|
|
| attn = torch.sigmoid(motion).unsqueeze(2) |
| attn = self.expand(attn) |
| attn = attn.expand(-1, -1, N, -1) |
| enhanced = x * attn |
| return x + self.delta.tanh() * enhanced |
|
|
|
|
| class Mlp(nn.Module): |
| def __init__(self, dim: int, mlp_ratio: float, drop: float): |
| super().__init__() |
| hidden_dim = int(dim * mlp_ratio) |
| self.fc1 = nn.Linear(dim, hidden_dim) |
| self.act = nn.GELU() |
| self.fc2 = nn.Linear(hidden_dim, dim) |
| self.drop = nn.Dropout(drop) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.fc1(x) |
| x = self.act(x) |
| x = self.drop(x) |
| x = self.fc2(x) |
| x = self.drop(x) |
| return x |
|
|
|
|
| class LSViTBlock(nn.Module): |
| def __init__(self, dim, num_heads, mlp_ratio, drop_rate, attn_drop, drop_path): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim) |
| self.attn = Attention(dim, num_heads, True, attn_drop, drop_rate) |
| self.drop_path1 = DropPath(drop_path) |
| self.norm2 = nn.LayerNorm(dim) |
| self.mlp = Mlp(dim, mlp_ratio, drop_rate) |
| self.drop_path2 = DropPath(drop_path) |
| self.lmim = LMIModule(dim) |
|
|
| def forward(self, x, B, T): |
| x = x + self.drop_path1(self.attn(self.norm1(x))) |
| x = x + self.drop_path2(self.mlp(self.norm2(x))) |
| BT, Np1, C = x.shape |
| assert BT == B * T |
| x = x.view(B, T, Np1, C) |
| x = self.lmim(x) |
| x = x.view(B * T, Np1, C) |
| return x |
|
|
|
|
| class LSViTBackbone(nn.Module): |
| def __init__(self, config: ViTConfig): |
| super().__init__() |
| self.config = config |
| self.patch_embed = PatchEmbed(config) |
| num_patches = self.patch_embed.num_patches |
|
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim)) |
| self.pos_drop = nn.Dropout(config.drop_rate) |
|
|
| dpr = torch.linspace(0, config.drop_path_rate, steps=config.depth).tolist() |
| self.blocks = nn.ModuleList( |
| [ |
| LSViTBlock( |
| dim=config.embed_dim, |
| num_heads=config.num_heads, |
| mlp_ratio=config.mlp_ratio, |
| drop_rate=config.drop_rate, |
| attn_drop=config.attn_drop_rate, |
| drop_path=dpr[i], |
| ) |
| for i in range(config.depth) |
| ] |
| ) |
|
|
| self.norm = nn.LayerNorm(config.embed_dim) |
|
|
| nn.init.trunc_normal_(self.cls_token, std=0.02) |
| nn.init.trunc_normal_(self.pos_embed, std=0.02) |
|
|
| def _interpolate_pos_encoding(self, x: torch.Tensor) -> torch.Tensor: |
| B, N, C = x.shape |
| num_patches = N - 1 |
| if num_patches == self.patch_embed.num_patches: |
| return self.pos_embed |
| cls_pos = self.pos_embed[:, :1] |
| patch_pos = self.pos_embed[:, 1:] |
| dim = patch_pos.shape[-1] |
| gs_old = int(math.sqrt(patch_pos.shape[1])) |
| gs_new = int(math.sqrt(num_patches)) |
| patch_pos = patch_pos.reshape(1, gs_old, gs_old, dim).permute(0, 3, 1, 2) |
| patch_pos = F.interpolate(patch_pos, size=(gs_new, gs_new), mode="bicubic", align_corners=False) |
| patch_pos = patch_pos.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, dim) |
| return torch.cat([cls_pos, patch_pos], dim=1) |
|
|
| def forward(self, video: torch.Tensor) -> torch.Tensor: |
| B, T, C, H, W = video.shape |
| x = video.reshape(B * T, C, H, W) |
| x = self.patch_embed(x) |
|
|
| cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) |
| x = torch.cat((cls_tokens, x), dim=1) |
|
|
| pos_embed = self._interpolate_pos_encoding(x) |
| x = x + pos_embed |
| x = self.pos_drop(x) |
|
|
| for block in self.blocks: |
| x = block(x, B, T) |
|
|
| x = self.norm(x) |
| x = x.view(B, T, x.shape[1], x.shape[2]) |
| return x |
|
|
|
|
| class LSViTForAction(nn.Module): |
| def __init__(self, config: ViTConfig, num_classes: int = 51, smif_window: int = 5): |
| super().__init__() |
| self.smif = SMIFModule(config.in_chans, window_size=smif_window) |
| self.backbone = LSViTBackbone(config) |
| self.head = nn.Linear(config.embed_dim, num_classes) |
|
|
| def forward(self, video: torch.Tensor) -> torch.Tensor: |
| x = self.smif(video) |
| feats = self.backbone(x) |
| cls_tokens = feats[:, :, 0] |
| pooled = cls_tokens.mean(dim=1) |
| logits = self.head(pooled) |
| return logits |
|
|
|
|
| HMDB51_CLASSES = [ |
| "brush_hair", "cartwheel", "catch", "chew", "clap", "climb", "climb_stairs", |
| "dive", "draw_sword", "dribble", "drink", "eat", "fall_floor", "fencing", |
| "flic_flac", "golf", "handstand", "hit", "hug", "jump", "kick", "kick_ball", |
| "kiss", "laugh", "pick", "pour", "pullup", "punch", "push", "pushup", |
| "ride_bike", "ride_horse", "run", "shake_hands", "shoot_ball", "shoot_bow", |
| "shoot_gun", "sit", "situp", "smile", "smoke", "somersault", "stand", |
| "swing_baseball", "sword", "sword_exercise", "talk", "throw", "turn", |
| "walk", "wave", |
| ] |
|
|
