Skip to content

zamba.models.slowfast_models

Classes

SlowFast

Bases: ZambaVideoClassificationLightningModule

Pretrained SlowFast model for fine-tuning with the following architecture:

Input -> SlowFast Base (including trainable Backbone) -> Res Basic Head -> Output

Attributes:

Name Type Description
backbone torch.nn.Module

When scheduling the backbone to train with the BackboneFinetune callback, this indicates the trainable part of the base.
base torch.nn.Module

The entire model prior to the head.
head torch.nn.Module

The trainable head.
_backbone_output_dim int

Dimensionality of the backbone output (and head input).
Source code in zamba/models/slowfast_models.py 
 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
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
@register_model
class SlowFast(ZambaVideoClassificationLightningModule):
    """Pretrained SlowFast model for fine-tuning with the following architecture:

    Input -> SlowFast Base (including trainable Backbone) -> Res Basic Head -> Output

    Attributes:
        backbone (torch.nn.Module): When scheduling the backbone to train with the
            `BackboneFinetune` callback, this indicates the trainable part of the base.
        base (torch.nn.Module): The entire model prior to the head.
        head (torch.nn.Module): The trainable head.
        _backbone_output_dim (int): Dimensionality of the backbone output (and head input).
    """

    def __init__(
        self,
        backbone_mode: str = "train",
        post_backbone_dropout: Optional[float] = None,
        output_with_global_average: bool = True,
        head_dropout_rate: Optional[float] = None,
        head_hidden_layer_sizes: Optional[Tuple[int]] = None,
        finetune_from: Optional[Union[os.PathLike, str]] = None,
        **kwargs,
    ):
        """Initializes the SlowFast model.

        Args:
            backbone_mode (str): If "eval", treat the backbone as a feature extractor
                and set to evaluation mode in all forward passes.
            post_backbone_dropout (float, optional): Dropout that operates on the output of the
                backbone + pool (before the fully-connected layer in the head).
            output_with_global_average (bool): If True, apply an adaptive average pooling
                operation after the fully-connected layer in the head.
            head_dropout_rate (float, optional): Optional dropout rate applied after backbone and
                between projection layers in the head.
            head_hidden_layer_sizes (tuple of int): If not None, the size of hidden layers in the
                head multilayer perceptron.
            finetune_from (pathlike or str, optional): If not None, load an existing model from
                the path and resume training from an existing model.
        """
        super().__init__(**kwargs)

        if finetune_from is None:
            self.initialize_from_torchub()
        else:
            model = self.load_from_checkpoint(finetune_from)
            self._backbone_output_dim = model.head.proj.in_features
            self.backbone = model.backbone
            self.base = model.base

        for param in self.base.parameters():
            param.requires_grad = False

        head = ResNetBasicHead(
            proj=build_multilayer_perceptron(
                self._backbone_output_dim,
                head_hidden_layer_sizes,
                self.num_classes,
                activation=torch.nn.ReLU,
                dropout=head_dropout_rate,
                output_activation=None,
            ),
            activation=None,
            pool=None,
            dropout=None
            if post_backbone_dropout is None
            else torch.nn.Dropout(post_backbone_dropout),
            output_pool=torch.nn.AdaptiveAvgPool3d(1),
        )

        self.backbone_mode = backbone_mode
        self.head = head

        self.save_hyperparameters(
            "backbone_mode",
            "head_dropout_rate",
            "head_hidden_layer_sizes",
            "output_with_global_average",
            "post_backbone_dropout",
        )

    def initialize_from_torchub(self):
        """Loads SlowFast model from torchhub and prepares ZambaVideoClassificationLightningModule
        by removing the head and setting the backbone and base."""

        # workaround for pytorch bug
        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
        base = torch.hub.load(
            "facebookresearch/pytorchvideo:0.1.3", model="slowfast_r50", pretrained=True
        )
        self._backbone_output_dim = base.blocks[-1].proj.in_features

        base.blocks = base.blocks[:-1]  # Remove the pre-trained head

        # self.backbone attribute lets `BackboneFinetune` freeze and unfreeze that module
        self.backbone = base.blocks[-2:]
        self.base = base

    def forward(self, x, *args, **kwargs):
        if self.backbone_mode == "eval":
            self.base.eval()

        x = self.base(x)
        return self.head(x)

Attributes

backbone = model.backbone instance-attribute
backbone_mode = backbone_mode instance-attribute
base = model.base instance-attribute
head = head instance-attribute

Functions

__init__(backbone_mode: str = 'train', post_backbone_dropout: Optional[float] = None, output_with_global_average: bool = True, head_dropout_rate: Optional[float] = None, head_hidden_layer_sizes: Optional[Tuple[int]] = None, finetune_from: Optional[Union[os.PathLike, str]] = None, **kwargs)

Initializes the SlowFast model.

Parameters:

Name Type Description Default
backbone_mode str

If "eval", treat the backbone as a feature extractor and set to evaluation mode in all forward passes.

 'train'
post_backbone_dropout float

Dropout that operates on the output of the backbone + pool (before the fully-connected layer in the head).

 None
output_with_global_average bool

If True, apply an adaptive average pooling operation after the fully-connected layer in the head.

 True
head_dropout_rate float

Optional dropout rate applied after backbone and between projection layers in the head.

 None
head_hidden_layer_sizes tuple of int

If not None, the size of hidden layers in the head multilayer perceptron.

 None
finetune_from pathlike or str

If not None, load an existing model from the path and resume training from an existing model.

 None
Source code in zamba/models/slowfast_models.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def __init__(
    self,
    backbone_mode: str = "train",
    post_backbone_dropout: Optional[float] = None,
    output_with_global_average: bool = True,
    head_dropout_rate: Optional[float] = None,
    head_hidden_layer_sizes: Optional[Tuple[int]] = None,
    finetune_from: Optional[Union[os.PathLike, str]] = None,
    **kwargs,
):
    """Initializes the SlowFast model.

    Args:
        backbone_mode (str): If "eval", treat the backbone as a feature extractor
            and set to evaluation mode in all forward passes.
        post_backbone_dropout (float, optional): Dropout that operates on the output of the
            backbone + pool (before the fully-connected layer in the head).
        output_with_global_average (bool): If True, apply an adaptive average pooling
            operation after the fully-connected layer in the head.
        head_dropout_rate (float, optional): Optional dropout rate applied after backbone and
            between projection layers in the head.
        head_hidden_layer_sizes (tuple of int): If not None, the size of hidden layers in the
            head multilayer perceptron.
        finetune_from (pathlike or str, optional): If not None, load an existing model from
            the path and resume training from an existing model.
    """
    super().__init__(**kwargs)

    if finetune_from is None:
        self.initialize_from_torchub()
    else:
        model = self.load_from_checkpoint(finetune_from)
        self._backbone_output_dim = model.head.proj.in_features
        self.backbone = model.backbone
        self.base = model.base

    for param in self.base.parameters():
        param.requires_grad = False

    head = ResNetBasicHead(
        proj=build_multilayer_perceptron(
            self._backbone_output_dim,
            head_hidden_layer_sizes,
            self.num_classes,
            activation=torch.nn.ReLU,
            dropout=head_dropout_rate,
            output_activation=None,
        ),
        activation=None,
        pool=None,
        dropout=None
        if post_backbone_dropout is None
        else torch.nn.Dropout(post_backbone_dropout),
        output_pool=torch.nn.AdaptiveAvgPool3d(1),
    )

    self.backbone_mode = backbone_mode
    self.head = head

    self.save_hyperparameters(
        "backbone_mode",
        "head_dropout_rate",
        "head_hidden_layer_sizes",
        "output_with_global_average",
        "post_backbone_dropout",
    )
forward(x, *args, **kwargs)
Source code in zamba/models/slowfast_models.py 
109
110
111
112
113
114
def forward(self, x, *args, **kwargs):
    if self.backbone_mode == "eval":
        self.base.eval()

    x = self.base(x)
    return self.head(x)
initialize_from_torchub()

Loads SlowFast model from torchhub and prepares ZambaVideoClassificationLightningModule by removing the head and setting the backbone and base.

Source code in zamba/models/slowfast_models.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def initialize_from_torchub(self):
    """Loads SlowFast model from torchhub and prepares ZambaVideoClassificationLightningModule
    by removing the head and setting the backbone and base."""

    # workaround for pytorch bug
    torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
    base = torch.hub.load(
        "facebookresearch/pytorchvideo:0.1.3", model="slowfast_r50", pretrained=True
    )
    self._backbone_output_dim = base.blocks[-1].proj.in_features

    base.blocks = base.blocks[:-1]  # Remove the pre-trained head

    # self.backbone attribute lets `BackboneFinetune` freeze and unfreeze that module
    self.backbone = base.blocks[-2:]
    self.base = base

Functions