import math from dataclasses import dataclass from functools import partial from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import torch import torch.fx import torch.nn as nn from ...ops import MLP, StochasticDepth from ...transforms._presets import VideoClassification from ...utils import _log_api_usage_once from .._api import register_model, Weights, WeightsEnum from .._meta import _KINETICS400_CATEGORIES from .._utils import _ovewrite_named_param, handle_legacy_interface __all__ = [ "MViT", "MViT_V1_B_Weights", "mvit_v1_b", "MViT_V2_S_Weights", "mvit_v2_s", ] @dataclass class MSBlockConfig: num_heads: int input_channels: int output_channels: int kernel_q: List[int] kernel_kv: List[int] stride_q: List[int] stride_kv: List[int] def _prod(s: Sequence[int]) -> int: product = 1 for v in s: product *= v return product def _unsqueeze(x: torch.Tensor, target_dim: int, expand_dim: int) -> Tuple[torch.Tensor, int]: tensor_dim = x.dim() if tensor_dim == target_dim - 1: x = x.unsqueeze(expand_dim) elif tensor_dim != target_dim: raise ValueError(f"Unsupported input dimension {x.shape}") return x, tensor_dim def _squeeze(x: torch.Tensor, target_dim: int, expand_dim: int, tensor_dim: int) -> torch.Tensor: if tensor_dim == target_dim - 1: x = x.squeeze(expand_dim) return x torch.fx.wrap("_unsqueeze") torch.fx.wrap("_squeeze") class Pool(nn.Module): def __init__( self, pool: nn.Module, norm: Optional[nn.Module], activation: Optional[nn.Module] = None, norm_before_pool: bool = False, ) -> None: super().__init__() self.pool = pool layers = [] if norm is not None: layers.append(norm) if activation is not None: layers.append(activation) self.norm_act = nn.Sequential(*layers) if layers else None self.norm_before_pool = norm_before_pool def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: x, tensor_dim = _unsqueeze(x, 4, 1) # Separate the class token and reshape the input class_token, x = torch.tensor_split(x, indices=(1,), dim=2) x = x.transpose(2, 3) B, N, C = x.shape[:3] x = x.reshape((B * N, C) + thw).contiguous() # normalizing prior pooling is useful when we use BN which can be absorbed to speed up inference if self.norm_before_pool and self.norm_act is not None: x = self.norm_act(x) # apply the pool on the input and add back the token x = self.pool(x) T, H, W = x.shape[2:] x = x.reshape(B, N, C, -1).transpose(2, 3) x = torch.cat((class_token, x), dim=2) if not self.norm_before_pool and self.norm_act is not None: x = self.norm_act(x) x = _squeeze(x, 4, 1, tensor_dim) return x, (T, H, W) def _interpolate(embedding: torch.Tensor, d: int) -> torch.Tensor: if embedding.shape[0] == d: return embedding return ( nn.functional.interpolate( embedding.permute(1, 0).unsqueeze(0), size=d, mode="linear", ) .squeeze(0) .permute(1, 0) ) def _add_rel_pos( attn: torch.Tensor, q: torch.Tensor, q_thw: Tuple[int, int, int], k_thw: Tuple[int, int, int], rel_pos_h: torch.Tensor, rel_pos_w: torch.Tensor, rel_pos_t: torch.Tensor, ) -> torch.Tensor: # Modified code from: https://github.com/facebookresearch/SlowFast/commit/1aebd71a2efad823d52b827a3deaf15a56cf4932 q_t, q_h, q_w = q_thw k_t, k_h, k_w = k_thw dh = int(2 * max(q_h, k_h) - 1) dw = int(2 * max(q_w, k_w) - 1) dt = int(2 * max(q_t, k_t) - 1) # Scale up rel pos if shapes for q and k are different. q_h_ratio = max(k_h / q_h, 1.0) k_h_ratio = max(q_h / k_h, 1.0) dist_h = torch.arange(q_h)[:, None] * q_h_ratio - (torch.arange(k_h)[None, :] + (1.0 - k_h)) * k_h_ratio q_w_ratio = max(k_w / q_w, 1.0) k_w_ratio = max(q_w / k_w, 1.0) dist_w = torch.arange(q_w)[:, None] * q_w_ratio - (torch.arange(k_w)[None, :] + (1.0 - k_w)) * k_w_ratio q_t_ratio = max(k_t / q_t, 1.0) k_t_ratio = max(q_t / k_t, 1.0) dist_t = torch.arange(q_t)[:, None] * q_t_ratio - (torch.arange(k_t)[None, :] + (1.0 - k_t)) * k_t_ratio # Interpolate rel pos if needed. rel_pos_h = _interpolate(rel_pos_h, dh) rel_pos_w = _interpolate(rel_pos_w, dw) rel_pos_t = _interpolate(rel_pos_t, dt) Rh = rel_pos_h[dist_h.long()] Rw = rel_pos_w[dist_w.long()] Rt = rel_pos_t[dist_t.long()] B, n_head, _, dim = q.shape r_q = q[:, :, 1:].reshape(B, n_head, q_t, q_h, q_w, dim) rel_h_q = torch.einsum("bythwc,hkc->bythwk", r_q, Rh) # [B, H, q_t, qh, qw, k_h] rel_w_q = torch.einsum("bythwc,wkc->bythwk", r_q, Rw) # [B, H, q_t, qh, qw, k_w] # [B, H, q_t, q_h, q_w, dim] -> [q_t, B, H, q_h, q_w, dim] -> [q_t, B*H*q_h*q_w, dim] r_q = r_q.permute(2, 0, 1, 3, 4, 5).reshape(q_t, B * n_head * q_h * q_w, dim) # [q_t, B*H*q_h*q_w, dim] * [q_t, dim, k_t] = [q_t, B*H*q_h*q_w, k_t] -> [B*H*q_h*q_w, q_t, k_t] rel_q_t = torch.matmul(r_q, Rt.transpose(1, 2)).transpose(0, 1) # [B*H*q_h*q_w, q_t, k_t] -> [B, H, q_t, q_h, q_w, k_t] rel_q_t = rel_q_t.view(B, n_head, q_h, q_w, q_t, k_t).permute(0, 1, 4, 2, 3, 5) # Combine rel pos. rel_pos = ( rel_h_q[:, :, :, :, :, None, :, None] + rel_w_q[:, :, :, :, :, None, None, :] + rel_q_t[:, :, :, :, :, :, None, None] ).reshape(B, n_head, q_t * q_h * q_w, k_t * k_h * k_w) # Add it to attention attn[:, :, 1:, 1:] += rel_pos return attn def _add_shortcut(x: torch.Tensor, shortcut: torch.Tensor, residual_with_cls_embed: bool): if residual_with_cls_embed: x.add_(shortcut) else: x[:, :, 1:, :] += shortcut[:, :, 1:, :] return x torch.fx.wrap("_add_rel_pos") torch.fx.wrap("_add_shortcut") class MultiscaleAttention(nn.Module): def __init__( self, input_size: List[int], embed_dim: int, output_dim: int, num_heads: int, kernel_q: List[int], kernel_kv: List[int], stride_q: List[int], stride_kv: List[int], residual_pool: bool, residual_with_cls_embed: bool, rel_pos_embed: bool, dropout: float = 0.0, norm_layer: Callable[..., nn.Module] = nn.LayerNorm, ) -> None: super().__init__() self.embed_dim = embed_dim self.output_dim = output_dim self.num_heads = num_heads self.head_dim = output_dim // num_heads self.scaler = 1.0 / math.sqrt(self.head_dim) self.residual_pool = residual_pool self.residual_with_cls_embed = residual_with_cls_embed self.qkv = nn.Linear(embed_dim, 3 * output_dim) layers: List[nn.Module] = [nn.Linear(output_dim, output_dim)] if dropout > 0.0: layers.append(nn.Dropout(dropout, inplace=True)) self.project = nn.Sequential(*layers) self.pool_q: Optional[nn.Module] = None if _prod(kernel_q) > 1 or _prod(stride_q) > 1: padding_q = [int(q // 2) for q in kernel_q] self.pool_q = Pool( nn.Conv3d( self.head_dim, self.head_dim, kernel_q, # type: ignore[arg-type] stride=stride_q, # type: ignore[arg-type] padding=padding_q, # type: ignore[arg-type] groups=self.head_dim, bias=False, ), norm_layer(self.head_dim), ) self.pool_k: Optional[nn.Module] = None self.pool_v: Optional[nn.Module] = None if _prod(kernel_kv) > 1 or _prod(stride_kv) > 1: padding_kv = [int(kv // 2) for kv in kernel_kv] self.pool_k = Pool( nn.Conv3d( self.head_dim, self.head_dim, kernel_kv, # type: ignore[arg-type] stride=stride_kv, # type: ignore[arg-type] padding=padding_kv, # type: ignore[arg-type] groups=self.head_dim, bias=False, ), norm_layer(self.head_dim), ) self.pool_v = Pool( nn.Conv3d( self.head_dim, self.head_dim, kernel_kv, # type: ignore[arg-type] stride=stride_kv, # type: ignore[arg-type] padding=padding_kv, # type: ignore[arg-type] groups=self.head_dim, bias=False, ), norm_layer(self.head_dim), ) self.rel_pos_h: Optional[nn.Parameter] = None self.rel_pos_w: Optional[nn.Parameter] = None self.rel_pos_t: Optional[nn.Parameter] = None if rel_pos_embed: size = max(input_size[1:]) q_size = size // stride_q[1] if len(stride_q) > 0 else size kv_size = size // stride_kv[1] if len(stride_kv) > 0 else size spatial_dim = 2 * max(q_size, kv_size) - 1 temporal_dim = 2 * input_size[0] - 1 self.rel_pos_h = nn.Parameter(torch.zeros(spatial_dim, self.head_dim)) self.rel_pos_w = nn.Parameter(torch.zeros(spatial_dim, self.head_dim)) self.rel_pos_t = nn.Parameter(torch.zeros(temporal_dim, self.head_dim)) nn.init.trunc_normal_(self.rel_pos_h, std=0.02) nn.init.trunc_normal_(self.rel_pos_w, std=0.02) nn.init.trunc_normal_(self.rel_pos_t, std=0.02) def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: B, N, C = x.shape q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).transpose(1, 3).unbind(dim=2) if self.pool_k is not None: k, k_thw = self.pool_k(k, thw) else: k_thw = thw if self.pool_v is not None: v = self.pool_v(v, thw)[0] if self.pool_q is not None: q, thw = self.pool_q(q, thw) attn = torch.matmul(self.scaler * q, k.transpose(2, 3)) if self.rel_pos_h is not None and self.rel_pos_w is not None and self.rel_pos_t is not None: attn = _add_rel_pos( attn, q, thw, k_thw, self.rel_pos_h, self.rel_pos_w, self.rel_pos_t, ) attn = attn.softmax(dim=-1) x = torch.matmul(attn, v) if self.residual_pool: _add_shortcut(x, q, self.residual_with_cls_embed) x = x.transpose(1, 2).reshape(B, -1, self.output_dim) x = self.project(x) return x, thw class MultiscaleBlock(nn.Module): def __init__( self, input_size: List[int], cnf: MSBlockConfig, residual_pool: bool, residual_with_cls_embed: bool, rel_pos_embed: bool, proj_after_attn: bool, dropout: float = 0.0, stochastic_depth_prob: float = 0.0, norm_layer: Callable[..., nn.Module] = nn.LayerNorm, ) -> None: super().__init__() self.proj_after_attn = proj_after_attn self.pool_skip: Optional[nn.Module] = None if _prod(cnf.stride_q) > 1: kernel_skip = [s + 1 if s > 1 else s for s in cnf.stride_q] padding_skip = [int(k // 2) for k in kernel_skip] self.pool_skip = Pool( nn.MaxPool3d(kernel_skip, stride=cnf.stride_q, padding=padding_skip), None # type: ignore[arg-type] ) attn_dim = cnf.output_channels if proj_after_attn else cnf.input_channels self.norm1 = norm_layer(cnf.input_channels) self.norm2 = norm_layer(attn_dim) self.needs_transposal = isinstance(self.norm1, nn.BatchNorm1d) self.attn = MultiscaleAttention( input_size, cnf.input_channels, attn_dim, cnf.num_heads, kernel_q=cnf.kernel_q, kernel_kv=cnf.kernel_kv, stride_q=cnf.stride_q, stride_kv=cnf.stride_kv, rel_pos_embed=rel_pos_embed, residual_pool=residual_pool, residual_with_cls_embed=residual_with_cls_embed, dropout=dropout, norm_layer=norm_layer, ) self.mlp = MLP( attn_dim, [4 * attn_dim, cnf.output_channels], activation_layer=nn.GELU, dropout=dropout, inplace=None, ) self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row") self.project: Optional[nn.Module] = None if cnf.input_channels != cnf.output_channels: self.project = nn.Linear(cnf.input_channels, cnf.output_channels) def forward(self, x: torch.Tensor, thw: Tuple[int, int, int]) -> Tuple[torch.Tensor, Tuple[int, int, int]]: x_norm1 = self.norm1(x.transpose(1, 2)).transpose(1, 2) if self.needs_transposal else self.norm1(x) x_attn, thw_new = self.attn(x_norm1, thw) x = x if self.project is None or not self.proj_after_attn else self.project(x_norm1) x_skip = x if self.pool_skip is None else self.pool_skip(x, thw)[0] x = x_skip + self.stochastic_depth(x_attn) x_norm2 = self.norm2(x.transpose(1, 2)).transpose(1, 2) if self.needs_transposal else self.norm2(x) x_proj = x if self.project is None or self.proj_after_attn else self.project(x_norm2) return x_proj + self.stochastic_depth(self.mlp(x_norm2)), thw_new class PositionalEncoding(nn.Module): def __init__(self, embed_size: int, spatial_size: Tuple[int, int], temporal_size: int, rel_pos_embed: bool) -> None: super().__init__() self.spatial_size = spatial_size self.temporal_size = temporal_size self.class_token = nn.Parameter(torch.zeros(embed_size)) self.spatial_pos: Optional[nn.Parameter] = None self.temporal_pos: Optional[nn.Parameter] = None self.class_pos: Optional[nn.Parameter] = None if not rel_pos_embed: self.spatial_pos = nn.Parameter(torch.zeros(self.spatial_size[0] * self.spatial_size[1], embed_size)) self.temporal_pos = nn.Parameter(torch.zeros(self.temporal_size, embed_size)) self.class_pos = nn.Parameter(torch.zeros(embed_size)) def forward(self, x: torch.Tensor) -> torch.Tensor: class_token = self.class_token.expand(x.size(0), -1).unsqueeze(1) x = torch.cat((class_token, x), dim=1) if self.spatial_pos is not None and self.temporal_pos is not None and self.class_pos is not None: hw_size, embed_size = self.spatial_pos.shape pos_embedding = torch.repeat_interleave(self.temporal_pos, hw_size, dim=0) pos_embedding.add_(self.spatial_pos.unsqueeze(0).expand(self.temporal_size, -1, -1).reshape(-1, embed_size)) pos_embedding = torch.cat((self.class_pos.unsqueeze(0), pos_embedding), dim=0).unsqueeze(0) x.add_(pos_embedding) return x class MViT(nn.Module): def __init__( self, spatial_size: Tuple[int, int], temporal_size: int, block_setting: Sequence[MSBlockConfig], residual_pool: bool, residual_with_cls_embed: bool, rel_pos_embed: bool, proj_after_attn: bool, dropout: float = 0.5, attention_dropout: float = 0.0, stochastic_depth_prob: float = 0.0, num_classes: int = 400, block: Optional[Callable[..., nn.Module]] = None, norm_layer: Optional[Callable[..., nn.Module]] = None, patch_embed_kernel: Tuple[int, int, int] = (3, 7, 7), patch_embed_stride: Tuple[int, int, int] = (2, 4, 4), patch_embed_padding: Tuple[int, int, int] = (1, 3, 3), ) -> None: """ MViT main class. Args: spatial_size (tuple of ints): The spacial size of the input as ``(H, W)``. temporal_size (int): The temporal size ``T`` of the input. block_setting (sequence of MSBlockConfig): The Network structure. residual_pool (bool): If True, use MViTv2 pooling residual connection. residual_with_cls_embed (bool): If True, the addition on the residual connection will include the class embedding. rel_pos_embed (bool): If True, use MViTv2's relative positional embeddings. proj_after_attn (bool): If True, apply the projection after the attention. dropout (float): Dropout rate. Default: 0.0. attention_dropout (float): Attention dropout rate. Default: 0.0. stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0. num_classes (int): The number of classes. block (callable, optional): Module specifying the layer which consists of the attention and mlp. norm_layer (callable, optional): Module specifying the normalization layer to use. patch_embed_kernel (tuple of ints): The kernel of the convolution that patchifies the input. patch_embed_stride (tuple of ints): The stride of the convolution that patchifies the input. patch_embed_padding (tuple of ints): The padding of the convolution that patchifies the input. """ super().__init__() # This implementation employs a different parameterization scheme than the one used at PyTorch Video: # https://github.com/facebookresearch/pytorchvideo/blob/718d0a4/pytorchvideo/models/vision_transformers.py # We remove any experimental configuration that didn't make it to the final variants of the models. To represent # the configuration of the architecture we use the simplified form suggested at Table 1 of the paper. _log_api_usage_once(self) total_stage_blocks = len(block_setting) if total_stage_blocks == 0: raise ValueError("The configuration parameter can't be empty.") if block is None: block = MultiscaleBlock if norm_layer is None: norm_layer = partial(nn.LayerNorm, eps=1e-6) # Patch Embedding module self.conv_proj = nn.Conv3d( in_channels=3, out_channels=block_setting[0].input_channels, kernel_size=patch_embed_kernel, stride=patch_embed_stride, padding=patch_embed_padding, ) input_size = [size // stride for size, stride in zip((temporal_size,) + spatial_size, self.conv_proj.stride)] # Spatio-Temporal Class Positional Encoding self.pos_encoding = PositionalEncoding( embed_size=block_setting[0].input_channels, spatial_size=(input_size[1], input_size[2]), temporal_size=input_size[0], rel_pos_embed=rel_pos_embed, ) # Encoder module self.blocks = nn.ModuleList() for stage_block_id, cnf in enumerate(block_setting): # adjust stochastic depth probability based on the depth of the stage block sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) self.blocks.append( block( input_size=input_size, cnf=cnf, residual_pool=residual_pool, residual_with_cls_embed=residual_with_cls_embed, rel_pos_embed=rel_pos_embed, proj_after_attn=proj_after_attn, dropout=attention_dropout, stochastic_depth_prob=sd_prob, norm_layer=norm_layer, ) ) if len(cnf.stride_q) > 0: input_size = [size // stride for size, stride in zip(input_size, cnf.stride_q)] self.norm = norm_layer(block_setting[-1].output_channels) # Classifier module self.head = nn.Sequential( nn.Dropout(dropout, inplace=True), nn.Linear(block_setting[-1].output_channels, num_classes), ) for m in self.modules(): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0.0) elif isinstance(m, nn.LayerNorm): if m.weight is not None: nn.init.constant_(m.weight, 1.0) if m.bias is not None: nn.init.constant_(m.bias, 0.0) elif isinstance(m, PositionalEncoding): for weights in m.parameters(): nn.init.trunc_normal_(weights, std=0.02) def forward(self, x: torch.Tensor) -> torch.Tensor: # Convert if necessary (B, C, H, W) -> (B, C, 1, H, W) x = _unsqueeze(x, 5, 2)[0] # patchify and reshape: (B, C, T, H, W) -> (B, embed_channels[0], T', H', W') -> (B, THW', embed_channels[0]) x = self.conv_proj(x) x = x.flatten(2).transpose(1, 2) # add positional encoding x = self.pos_encoding(x) # pass patches through the encoder thw = (self.pos_encoding.temporal_size,) + self.pos_encoding.spatial_size for block in self.blocks: x, thw = block(x, thw) x = self.norm(x) # classifier "token" as used by standard language architectures x = x[:, 0] x = self.head(x) return x def _mvit( block_setting: List[MSBlockConfig], stochastic_depth_prob: float, weights: Optional[WeightsEnum], progress: bool, **kwargs: Any, ) -> MViT: if weights is not None: _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) assert weights.meta["min_size"][0] == weights.meta["min_size"][1] _ovewrite_named_param(kwargs, "spatial_size", weights.meta["min_size"]) _ovewrite_named_param(kwargs, "temporal_size", weights.meta["min_temporal_size"]) spatial_size = kwargs.pop("spatial_size", (224, 224)) temporal_size = kwargs.pop("temporal_size", 16) model = MViT( spatial_size=spatial_size, temporal_size=temporal_size, block_setting=block_setting, residual_pool=kwargs.pop("residual_pool", False), residual_with_cls_embed=kwargs.pop("residual_with_cls_embed", True), rel_pos_embed=kwargs.pop("rel_pos_embed", False), proj_after_attn=kwargs.pop("proj_after_attn", False), stochastic_depth_prob=stochastic_depth_prob, **kwargs, ) if weights is not None: model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) return model class MViT_V1_B_Weights(WeightsEnum): KINETICS400_V1 = Weights( url="https://download.pytorch.org/models/mvit_v1_b-dbeb1030.pth", transforms=partial( VideoClassification, crop_size=(224, 224), resize_size=(256,), mean=(0.45, 0.45, 0.45), std=(0.225, 0.225, 0.225), ), meta={ "min_size": (224, 224), "min_temporal_size": 16, "categories": _KINETICS400_CATEGORIES, "recipe": "https://github.com/facebookresearch/pytorchvideo/blob/main/docs/source/model_zoo.md", "_docs": ( "The weights were ported from the paper. The accuracies are estimated on video-level " "with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`" ), "num_params": 36610672, "_metrics": { "Kinetics-400": { "acc@1": 78.477, "acc@5": 93.582, } }, "_ops": 70.599, "_file_size": 139.764, }, ) DEFAULT = KINETICS400_V1 class MViT_V2_S_Weights(WeightsEnum): KINETICS400_V1 = Weights( url="https://download.pytorch.org/models/mvit_v2_s-ae3be167.pth", transforms=partial( VideoClassification, crop_size=(224, 224), resize_size=(256,), mean=(0.45, 0.45, 0.45), std=(0.225, 0.225, 0.225), ), meta={ "min_size": (224, 224), "min_temporal_size": 16, "categories": _KINETICS400_CATEGORIES, "recipe": "https://github.com/facebookresearch/SlowFast/blob/main/MODEL_ZOO.md", "_docs": ( "The weights were ported from the paper. The accuracies are estimated on video-level " "with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`" ), "num_params": 34537744, "_metrics": { "Kinetics-400": { "acc@1": 80.757, "acc@5": 94.665, } }, "_ops": 64.224, "_file_size": 131.884, }, ) DEFAULT = KINETICS400_V1 @register_model() @handle_legacy_interface(weights=("pretrained", MViT_V1_B_Weights.KINETICS400_V1)) def mvit_v1_b(*, weights: Optional[MViT_V1_B_Weights] = None, progress: bool = True, **kwargs: Any) -> MViT: """ Constructs a base MViTV1 architecture from `Multiscale Vision Transformers `__. .. betastatus:: video module Args: weights (:class:`~torchvision.models.video.MViT_V1_B_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.video.MViT_V1_B_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.video.MViT`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.video.MViT_V1_B_Weights :members: """ weights = MViT_V1_B_Weights.verify(weights) config: Dict[str, List] = { "num_heads": [1, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 8, 8], "input_channels": [96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768], "output_channels": [192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768, 768], "kernel_q": [[], [3, 3, 3], [], [3, 3, 3], [], [], [], [], [], [], [], [], [], [], [3, 3, 3], []], "kernel_kv": [ [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], ], "stride_q": [[], [1, 2, 2], [], [1, 2, 2], [], [], [], [], [], [], [], [], [], [], [1, 2, 2], []], "stride_kv": [ [1, 8, 8], [1, 4, 4], [1, 4, 4], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 1, 1], [1, 1, 1], ], } block_setting = [] for i in range(len(config["num_heads"])): block_setting.append( MSBlockConfig( num_heads=config["num_heads"][i], input_channels=config["input_channels"][i], output_channels=config["output_channels"][i], kernel_q=config["kernel_q"][i], kernel_kv=config["kernel_kv"][i], stride_q=config["stride_q"][i], stride_kv=config["stride_kv"][i], ) ) return _mvit( spatial_size=(224, 224), temporal_size=16, block_setting=block_setting, residual_pool=False, residual_with_cls_embed=False, stochastic_depth_prob=kwargs.pop("stochastic_depth_prob", 0.2), weights=weights, progress=progress, **kwargs, ) @register_model() @handle_legacy_interface(weights=("pretrained", MViT_V2_S_Weights.KINETICS400_V1)) def mvit_v2_s(*, weights: Optional[MViT_V2_S_Weights] = None, progress: bool = True, **kwargs: Any) -> MViT: """Constructs a small MViTV2 architecture from `Multiscale Vision Transformers `__ and `MViTv2: Improved Multiscale Vision Transformers for Classification and Detection `__. .. betastatus:: video module Args: weights (:class:`~torchvision.models.video.MViT_V2_S_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.video.MViT_V2_S_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.video.MViT`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.video.MViT_V2_S_Weights :members: """ weights = MViT_V2_S_Weights.verify(weights) config: Dict[str, List] = { "num_heads": [1, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 8, 8], "input_channels": [96, 96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768], "output_channels": [96, 192, 192, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 384, 768, 768], "kernel_q": [ [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], ], "kernel_kv": [ [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], ], "stride_q": [ [1, 1, 1], [1, 2, 2], [1, 1, 1], [1, 2, 2], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 2, 2], [1, 1, 1], ], "stride_kv": [ [1, 8, 8], [1, 4, 4], [1, 4, 4], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 1, 1], [1, 1, 1], ], } block_setting = [] for i in range(len(config["num_heads"])): block_setting.append( MSBlockConfig( num_heads=config["num_heads"][i], input_channels=config["input_channels"][i], output_channels=config["output_channels"][i], kernel_q=config["kernel_q"][i], kernel_kv=config["kernel_kv"][i], stride_q=config["stride_q"][i], stride_kv=config["stride_kv"][i], ) ) return _mvit( spatial_size=(224, 224), temporal_size=16, block_setting=block_setting, residual_pool=True, residual_with_cls_embed=False, rel_pos_embed=True, proj_after_attn=True, stochastic_depth_prob=kwargs.pop("stochastic_depth_prob", 0.2), weights=weights, progress=progress, **kwargs, )