File size: 1,381 Bytes
63f8a47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """Abstract frozen video-backbone interface used by the V-PVP head.
A backbone must expose:
- ``extract_patch_features(x, layer=None)`` returning patch tokens of shape
``(B, P, T', D)``;
- its input spec (``img_size``, ``patch_size``, ``num_frames``,
``input_mean``, ``input_std``);
- ``max_extract_layer``, the highest 0-indexed block that can be tapped.
The V-PVP head is agnostic to P, T' and D.
"""
from abc import ABC, abstractmethod
import torch.nn as nn
class FrozenVideoBackbone(nn.Module, ABC):
# Subclasses override these (usually per instance in __init__).
img_size: int = 224
patch_size: int = 16
num_frames: int = 16
num_patches_spatial: int = 196
embed_dim: int = 768
input_mean: tuple = (0.485, 0.456, 0.406)
input_std: tuple = (0.229, 0.224, 0.225)
max_extract_layer: int = 11
@abstractmethod
def extract_patch_features(self, x, layer=None):
"""x: (B, C, T, H, W), channel-normalized with (input_mean, input_std).
Returns (B, P, T', D), where T' is the backbone's output temporal
length (T // 2 for VideoMAE with tubelet_size=2). With ``layer=None``
the tokens are the output of the last transformer block, taken before
any final normalization. Otherwise block ``layer`` (0-indexed, must be
<= ``max_extract_layer``) is tapped.
"""
...
|