from functools import partial from typing import Any, Callable, Optional import torch from torch import nn from torchvision.ops.misc import Conv3dNormActivation 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__ = [ "S3D", "S3D_Weights", "s3d", ] class TemporalSeparableConv(nn.Sequential): def __init__( self, in_planes: int, out_planes: int, kernel_size: int, stride: int, padding: int, norm_layer: Callable[..., nn.Module], ): super().__init__( Conv3dNormActivation( in_planes, out_planes, kernel_size=(1, kernel_size, kernel_size), stride=(1, stride, stride), padding=(0, padding, padding), bias=False, norm_layer=norm_layer, ), Conv3dNormActivation( out_planes, out_planes, kernel_size=(kernel_size, 1, 1), stride=(stride, 1, 1), padding=(padding, 0, 0), bias=False, norm_layer=norm_layer, ), ) class SepInceptionBlock3D(nn.Module): def __init__( self, in_planes: int, b0_out: int, b1_mid: int, b1_out: int, b2_mid: int, b2_out: int, b3_out: int, norm_layer: Callable[..., nn.Module], ): super().__init__() self.branch0 = Conv3dNormActivation(in_planes, b0_out, kernel_size=1, stride=1, norm_layer=norm_layer) self.branch1 = nn.Sequential( Conv3dNormActivation(in_planes, b1_mid, kernel_size=1, stride=1, norm_layer=norm_layer), TemporalSeparableConv(b1_mid, b1_out, kernel_size=3, stride=1, padding=1, norm_layer=norm_layer), ) self.branch2 = nn.Sequential( Conv3dNormActivation(in_planes, b2_mid, kernel_size=1, stride=1, norm_layer=norm_layer), TemporalSeparableConv(b2_mid, b2_out, kernel_size=3, stride=1, padding=1, norm_layer=norm_layer), ) self.branch3 = nn.Sequential( nn.MaxPool3d(kernel_size=(3, 3, 3), stride=1, padding=1), Conv3dNormActivation(in_planes, b3_out, kernel_size=1, stride=1, norm_layer=norm_layer), ) def forward(self, x): x0 = self.branch0(x) x1 = self.branch1(x) x2 = self.branch2(x) x3 = self.branch3(x) out = torch.cat((x0, x1, x2, x3), 1) return out class S3D(nn.Module): """S3D main class. Args: num_class (int): number of classes for the classification task. dropout (float): dropout probability. norm_layer (Optional[Callable]): Module specifying the normalization layer to use. Inputs: x (Tensor): batch of videos with dimensions (batch, channel, time, height, width) """ def __init__( self, num_classes: int = 400, dropout: float = 0.2, norm_layer: Optional[Callable[..., torch.nn.Module]] = None, ) -> None: super().__init__() _log_api_usage_once(self) if norm_layer is None: norm_layer = partial(nn.BatchNorm3d, eps=0.001, momentum=0.001) self.features = nn.Sequential( TemporalSeparableConv(3, 64, 7, 2, 3, norm_layer), nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)), Conv3dNormActivation( 64, 64, kernel_size=1, stride=1, norm_layer=norm_layer, ), TemporalSeparableConv(64, 192, 3, 1, 1, norm_layer), nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)), SepInceptionBlock3D(192, 64, 96, 128, 16, 32, 32, norm_layer), SepInceptionBlock3D(256, 128, 128, 192, 32, 96, 64, norm_layer), nn.MaxPool3d(kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)), SepInceptionBlock3D(480, 192, 96, 208, 16, 48, 64, norm_layer), SepInceptionBlock3D(512, 160, 112, 224, 24, 64, 64, norm_layer), SepInceptionBlock3D(512, 128, 128, 256, 24, 64, 64, norm_layer), SepInceptionBlock3D(512, 112, 144, 288, 32, 64, 64, norm_layer), SepInceptionBlock3D(528, 256, 160, 320, 32, 128, 128, norm_layer), nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 0, 0)), SepInceptionBlock3D(832, 256, 160, 320, 32, 128, 128, norm_layer), SepInceptionBlock3D(832, 384, 192, 384, 48, 128, 128, norm_layer), ) self.avgpool = nn.AvgPool3d(kernel_size=(2, 7, 7), stride=1) self.classifier = nn.Sequential( nn.Dropout(p=dropout), nn.Conv3d(1024, num_classes, kernel_size=1, stride=1, bias=True), ) def forward(self, x): x = self.features(x) x = self.avgpool(x) x = self.classifier(x) x = torch.mean(x, dim=(2, 3, 4)) return x class S3D_Weights(WeightsEnum): KINETICS400_V1 = Weights( url="https://download.pytorch.org/models/s3d-d76dad2f.pth", transforms=partial( VideoClassification, crop_size=(224, 224), resize_size=(256, 256), ), meta={ "min_size": (224, 224), "min_temporal_size": 14, "categories": _KINETICS400_CATEGORIES, "recipe": "https://github.com/pytorch/vision/tree/main/references/video_classification#s3d", "_docs": ( "The weights aim to approximate the accuracy of the paper. The accuracies are estimated on clip-level " "with parameters `frame_rate=15`, `clips_per_video=1`, and `clip_len=128`." ), "num_params": 8320048, "_metrics": { "Kinetics-400": { "acc@1": 68.368, "acc@5": 88.050, } }, "_ops": 17.979, "_file_size": 31.972, }, ) DEFAULT = KINETICS400_V1 @register_model() @handle_legacy_interface(weights=("pretrained", S3D_Weights.KINETICS400_V1)) def s3d(*, weights: Optional[S3D_Weights] = None, progress: bool = True, **kwargs: Any) -> S3D: """Construct Separable 3D CNN model. Reference: `Rethinking Spatiotemporal Feature Learning `__. .. betastatus:: video module Args: weights (:class:`~torchvision.models.video.S3D_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.video.S3D_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.video.S3D`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.video.S3D_Weights :members: """ weights = S3D_Weights.verify(weights) if weights is not None: _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"])) model = S3D(**kwargs) if weights is not None: model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) return model