a

__init__.py

@@ -0,0 +1,74 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all.
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+# rely on isort to merge the imports
+from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
+
+
+_import_structure = {"configuration_omnivore": ["OMNIVORE_PRETRAINED_CONFIG_ARCHIVE_MAP", "OmnivoreConfig"]}
+
+try:
+    if not is_vision_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["feature_extraction_omnivore"] = ["OmnivoreFeatureExtractor"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_omnivore"] = [
+        "OMNIVORE_PRETRAINED_MODEL_ARCHIVE_LIST",
+        "OmnivoreForJointClassification",
+        "OmnivoreModel",
+        "OmnivorePreTrainedModel",
+    ]
+
+if TYPE_CHECKING:
+    from .configuration_omnivore import OMNIVORE_PRETRAINED_CONFIG_ARCHIVE_MAP, OmnivoreConfig
+
+    try:
+        if not is_vision_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .feature_extraction_omnivore import OmnivoreFeatureExtractor
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_omnivore import (
+            OMNIVORE_PRETRAINED_MODEL_ARCHIVE_LIST,
+            OmnivoreForClassification,
+            OmnivoreModel,
+            OmnivorePreTrainedModel,
+        )
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)

configuration_omnivore.py

@@ -0,0 +1,128 @@
+# coding=utf-8
+# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Omnivore model configuration"""
+
+from torch import nn
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+OMNIVORE_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "anugunj/omnivore": "https://huggingface.co/anugunj/omnivore/resolve/main/config.json",
+}
+
+
+class OmnivoreConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`OmnivoreModel`]. It is used to instantiate an
+    Omnivore model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the Omnivore
+    [anugunj/omnivore](https://huggingface.co/anugunj/omnivore) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        input_channels (`int`, *optional*, defaults to 3):
+            The number of input channels.
+        patch_size (`int` | `List[int]`, *optional*, defaults to [4, 4, 4]):
+            Patch size to use in the patch embedding layer.
+        embed_dim (`int`, *optional*, defaults to 96):
+            Number of linear projection output channels.
+        depths (`List[int]`, *optional*, defaults to [2, 2, 6, 2],):
+            Depth (number of layers) for each stage.
+        num_heads (`List[int]`, *optional*, defaults to [3, 6, 12, 24]):
+            Number of attention head of each stage.
+        window_size (`int`, *optional*, defaults to 7)
+            Size of the window used by swin transformer in the model,
+        mlp_ratios (`float`, *optional*, defaults to 4.0):
+            Ratio of the size of the hidden layer compared to the size of the input layer of the Mix FFNs in the
+            encoder blocks.
+        attention_dropout_rate (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        dropout_rate (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the patch embeddings probabilities and projections in attention.
+        drop_path_rate (`List[float]`, *optional*, defaults to `[0.0, 0.0, 0.1]`):
+            The dropout probability for stochastic depth, used in the blocks of the Transformer encoder.
+        qkv_bias (`bool`, *optional*, defaults to True):
+            The bias bool for query, key and value in attentions
+        qk_scale (`bool`, *optional*, defaults to None):
+            Override default qk scale of head_dim ** -0.5 if set.
+        norm_layer (`nn.Module`, *optional*, defaults to nn.LayerNorm):
+            Normalization layer for the model
+        patch_norm (`bool`, *optional*, defaults to False):
+            If True, add normalization after patch embedding.
+        frozen_stages (`int`, *optional*, defaults to -1):
+            Stages to be frozen (stop grad and set eval mode) -1 means not freezing any parameters.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+
+    Example:
+    ```python
+    >>> from transformers import OmnivoreModel, OmnivoreConfig
+
+    >>> # Initializing a Omnivore omnivore-tiny-224 style configuration
+    >>> configuration = OmnivoreConfig()
+    >>> # Initializing a model from the omnivore-tiny-224 style configuration
+    >>> model = OmnivoreModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+    model_type = "omnivore"
+
+    def __init__(
+        self,
+        input_channels=3,
+        patch_size=[2, 4, 4],
+        embed_dim=96,
+        depths=[2, 2, 18, 2],
+        num_heads=[3, 6, 12, 24],
+        window_size=(8, 7, 7),
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        dropout_rate=0.0,
+        attention_dropout_rate=0.0,
+        drop_path_rate=0.3,
+        patch_norm=True,
+        frozen_stages=-1,
+        depth_mode="summed_rgb_d_tokens",
+        initializer_range=0.02,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.input_channels = input_channels
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.depths = depths
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.mlp_ratio = mlp_ratio
+        self.qkv_bias = qkv_bias
+        self.qk_scale = qk_scale
+        self.dropout_rate = dropout_rate
+        self.attention_dropout_rate = attention_dropout_rate
+        self.drop_path_rate = drop_path_rate
+        self.patch_norm = patch_norm
+        self.frozen_stages = frozen_stages
+        self.initializer_range = initializer_range
+        self.head_dim_in = embed_dim * 8
+        self.depth_mode = depth_mode
+        self.num_image_labels = 1000
+        self.num_video_labels = 400
+        self.num_rgbd_labels = 19

modelling.py

@@ -0,0 +1,1130 @@
+# coding=utf-8
+# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch Omnivore model."""
+
+import math
+import warnings
+from dataclasses import dataclass
+from functools import lru_cache, reduce
+from operator import mul
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.utils.checkpoint
+import torch.utils.checkpoint as checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn import functional as F
+
+from transformers.utils.generic import ModelOutput
+
+from ...activations import ACT2FN
+from ...modeling_outputs import BaseModelOutputWithNoAttention, BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
+from ...modeling_utils import PreTrainedModel
+from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
+from .configuration_omnivore import OmnivoreConfig
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "OmnivoreConfig"
+_FEAT_EXTRACTOR_FOR_DOC = "OmniverseFeatureExtractor"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "anugunj/omnivore"
+_EXPECTED_OUTPUT_SHAPE = [1, 768, 7, 7]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "anugunj/omnivore"
+_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
+
+OMNIVORE_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "anugunj/omnivore",
+    # See all Omnivore models at https://huggingface.co/models?filter=omnivore
+]
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    r"""Fills the input Tensor with values drawn from a truncated
+    Args:
+    normal distribution. The values are effectively drawn from the normal distribution :math:`\mathcal{N}(\text{mean},
+    \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for
+    generating the random values works best when :math:`a \leq \text{mean} \leq b`.
+        tensor: an n-dimensional `torch.Tensor` mean: the mean of the normal distribution std: the standard deviation
+        of the normal distribution a: the minimum cutoff value b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5) >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+# Stochastic depth implementation
+# Taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the
+    DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as 'Drop
+    Connect' is a different form of dropout in a separate paper... See discussion:
+    https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the layer and
+    argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class OmnivoreDropPath(nn.Module):
+    def __init__(self, drop_prob=None):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class OmnivoreLayerNorm(nn.Module):
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError(f"Unsupported data format: {self.data_format}")
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor):
+        if self.data_format == "channels_last":
+            x = torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+class OmnivoreIm2Video(nn.Module):
+    """Convert Image into a trivial video"""
+
+    def forward(self, pixel_values):
+        if pixel_values.ndim == 4:
+            return pixel_values.unsqueeze(2)
+        elif pixel_values.ndim == 5:
+            return pixel_values
+        else:
+            raise ValueError(f"Dimension incorrect {pixel_values.shape}")
+
+
+class OmnivoreMLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, dropout_rate=0.0, act_layer=nn.GELU):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.linear1 = nn.Linear(in_features, hidden_features)
+        self.activation = act_layer()
+        self.linear2 = nn.Linear(hidden_features, out_features)
+        self.drop_out = nn.Dropout(dropout_rate)
+
+    def forward(self, hidden_state):
+        hidden_state = self.linear1(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        hidden_state = self.drop_out(hidden_state)
+        hidden_state = self.linear2(hidden_state)
+        hidden_state = self.drop_out(hidden_state)
+        return hidden_state
+
+
+def window_partition(input_feature, window_size):
+    batch_size, D, height, width, channels = input_feature.shape
+    input_feature = input_feature.view(
+        batch_size,
+        D // window_size[0],
+        window_size[0],
+        height // window_size[1],
+        window_size[1],
+        width // window_size[2],
+        window_size[2],
+        channels,
+    )
+    windows = input_feature.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, reduce(mul, window_size), channels)
+    return windows
+
+
+def window_partition_image(input_feature, window_size):
+    batch_size, height, width, channels = input_feature.shape
+    input_feature = input_feature.view(
+        batch_size, height // window_size[1], window_size[1], width // window_size[2], window_size[2], channels
+    )
+    windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[1], window_size[2], channels)
+    return windows
+
+
+def window_reverse(windows, windows_size, batch_size, D, height, width):
+    input_feature = windows.view(
+        batch_size,
+        D // windows_size[0],
+        height // windows_size[1],
+        width // windows_size[2],
+        windows_size[0],
+        windows_size[1],
+        windows_size[2],
+        -1,
+    )
+    input_feature = input_feature.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(batch_size, D, height, width, -1)
+    return input_feature
+
+
+def get_window_size(input_size, window_size, shift_size=None):
+    use_window_size = list(window_size)
+    if shift_size is not None:
+        use_shift_size = list(shift_size)
+    for i in range(len(input_size)):
+        if input_size[i] <= window_size[i]:
+            use_window_size[i] = input_size[i]
+            if shift_size is not None:
+                use_shift_size[i] = 0
+
+    if shift_size is None:
+        return tuple(use_window_size)
+    else:
+        return tuple(use_window_size), tuple(use_shift_size)
+
+
+class OmnivoreWindowAttention3D(nn.Module):
+    def __init__(
+        self,
+        dim,
+        window_size,
+        num_heads,
+        qkv_bias=False,
+        qk_scale=None,
+        attention_dropout_rate=0.0,
+        projection_dropout_rate=0.0,
+    ):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(
+                (2 * window_size[0] - 1) * (2 * window_size[1] - 1) * (2 * window_size[2] - 1),
+                num_heads,
+            )
+        )
+
+        # get pair-wise relative position index for each token inside the window
+        coords_d = torch.arange(self.window_size[0])
+        coords_h = torch.arange(self.window_size[1])
+        coords_w = torch.arange(self.window_size[2])
+        coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w))
+        coords_flatten = torch.flatten(coords, 1)
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 2] += self.window_size[2] - 1
+
+        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
+        relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1
+        relative_position_index = relative_coords.sum(-1)
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.queries_keys_values = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attention_dropout = nn.Dropout(attention_dropout_rate)
+        self.projection = nn.Linear(dim, dim)
+        self.projection_dropout = nn.Dropout(projection_dropout_rate)
+
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, hidden_state, attention_mask=None):
+        batch_size, seq_len, channels = hidden_state.shape
+        queries_keys_values = (
+            self.queries_keys_values(hidden_state)
+            .reshape(batch_size, seq_len, 3, self.num_heads, channels // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        queries, keys, values = queries_keys_values[0], queries_keys_values[1], queries_keys_values[2]
+
+        queries = queries * self.scale
+        attention = queries @ keys.transpose(-2, -1)
+
+        relative_position_bias = self.relative_position_bias_table[
+            self.relative_position_index[:seq_len, :seq_len].reshape(-1)
+        ].reshape(seq_len, seq_len, -1)
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        attention = attention + relative_position_bias.unsqueeze(0)
+
+        if attention_mask is not None:
+            nW = attention_mask.shape[0]
+            attention = attention.view(
+                batch_size // nW, nW, self.num_heads, seq_len, seq_len
+            ) + attention_mask.unsqueeze(1).unsqueeze(0)
+            attention = attention.view(-1, self.num_heads, seq_len, seq_len)
+            attention = self.softmax(attention)
+        else:
+            attention = self.softmax(attention)
+
+        attention = self.attention_dropout(attention)
+
+        hidden_state = (attention @ values).transpose(1, 2).reshape(batch_size, seq_len, channels)
+        hidden_state = self.projection(hidden_state)
+        hidden_state = self.projection_dropout(hidden_state)
+        return hidden_state
+
+
+class OmnivoreSwinTransformer3DLayer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(2, 7, 7),
+        shift_size=(0, 0, 0),
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        dropout_rate=0.0,
+        attention_dropout_rate=0.0,
+        drop_path_rate=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+
+        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-window_size"
+        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-window_size"
+        assert 0 <= self.shift_size[2] < self.window_size[2], "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attention = OmnivoreWindowAttention3D(
+            dim,
+            window_size=self.window_size,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attention_dropout_rate=attention_dropout_rate,
+            projection_dropout_rate=dropout_rate,
+        )
+
+        self.drop_path = OmnivoreDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = OmnivoreMLP(
+            in_features=dim, hidden_features=mlp_hidden_dim, dropout_rate=dropout_rate, act_layer=act_layer
+        )
+
+    def forward_before(self, hidden_state, attention_mask):
+        batch_size, D, height, width, channels = hidden_state.shape
+        window_size, shift_size = get_window_size((D, height, width), self.window_size, self.shift_size)
+
+        hidden_state = self.norm1(hidden_state)
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = pad_d0 = 0
+        pad_d1 = (window_size[0] - D % window_size[0]) % window_size[0]
+        pad_b = (window_size[1] - height % window_size[1]) % window_size[1]
+        pad_r = (window_size[2] - width % window_size[2]) % window_size[2]
+        hidden_state = F.pad(hidden_state, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1))
+        _, Dp, Hp, Wp, _ = hidden_state.shape
+        # cyclic shift
+        if any(i > 0 for i in shift_size):
+            shifted_hidden_state = torch.roll(
+                hidden_state, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3)
+            )
+            attention_mask = attention_mask
+        else:
+            shifted_hidden_state = hidden_state
+            attention_mask = None
+        # partition windows
+        hidden_state_windows = window_partition(shifted_hidden_state, window_size)
+        # W-MSA/SW-MSA
+        attention_windows = self.attention(hidden_state_windows, attention_mask=attention_mask)
+        # merge windows
+        attention_windows = attention_windows.view(-1, *(window_size + (channels,)))
+        shifted_hidden_state = window_reverse(attention_windows, window_size, batch_size, Dp, Hp, Wp)
+        # reverse cyclic shift
+        if any(i > 0 for i in shift_size):
+            hidden_state = torch.roll(
+                shifted_hidden_state, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3)
+            )
+        else:
+            hidden_state = shifted_hidden_state
+
+        if pad_d1 > 0 or pad_r > 0 or pad_b > 0:
+            hidden_state = hidden_state[:, :D, :height, :width, :].contiguous()
+        return hidden_state
+
+    def forward_after(self, hidden_state):
+        hidden_state = self.norm2(hidden_state)
+        hidden_state = self.mlp(hidden_state)
+        hidden_state = self.drop_path(hidden_state)
+        return hidden_state
+
+    def forward(self, hidden_state, mask_matrix, use_checkpoint=False):
+        shortcut = hidden_state
+        if use_checkpoint:
+            hidden_state = checkpoint.checkpoint(self.forward_before, hidden_state, mask_matrix)
+        else:
+            hidden_state = self.forward_before(hidden_state, mask_matrix)
+        hidden_state = shortcut + self.drop_path(hidden_state)
+
+        if use_checkpoint:
+            hidden_state = hidden_state + checkpoint.checkpoint(self.forward_after, hidden_state)
+        else:
+            hidden_state = hidden_state + self.forward_after(hidden_state)
+
+        return hidden_state
+
+
+class OmnivorePatchMerging(nn.Module):
+    """
+    Args:
+    Patch Merging Layer
+        dim (`int`): Number of input channels. norm_layer (`nn.Module`, *optional*): Normalization layer. Default:
+        `nn.LayerNorm`
+    """
+
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, hidden_state, height=None, width=None):
+        if height is None:
+            batch_size, D, height, width, channels = hidden_state.shape
+
+        # padding
+        pad_input = (height % 2 == 1) or (width % 2 == 1)
+        if pad_input:
+            hidden_state = F.pad(hidden_state, (0, 0, 0, width % 2, 0, height % 2))
+
+        hidden_state0 = hidden_state[:, :, 0::2, 0::2, :]
+        hidden_state1 = hidden_state[:, :, 1::2, 0::2, :]
+        hidden_state2 = hidden_state[:, :, 0::2, 1::2, :]
+        hidden_state3 = hidden_state[:, :, 1::2, 1::2, :]
+        hidden_state = torch.cat([hidden_state0, hidden_state1, hidden_state2, hidden_state3], -1)
+
+        hidden_state = self.norm(hidden_state)
+        hidden_state = self.reduction(hidden_state)
+
+        return hidden_state
+
+
+@lru_cache()
+def compute_mask(D, height, width, window_size, shift_size, device):
+    img_mask = torch.zeros((1, D, height, width, 1), device=device)  # 1 Dp Hp Wp 1
+    cnt = 0
+    for d in (
+        slice(-window_size[0]),
+        slice(-window_size[0], -shift_size[0]),
+        slice(-shift_size[0], None),
+    ):
+        for h in (
+            slice(-window_size[1]),
+            slice(-window_size[1], -shift_size[1]),
+            slice(-shift_size[1], None),
+        ):
+            for w in (
+                slice(-window_size[2]),
+                slice(-window_size[2], -shift_size[2]),
+                slice(-shift_size[2], None),
+            ):
+                img_mask[:, d, h, w, :] = cnt
+                cnt += 1
+    mask_windows = window_partition(img_mask, window_size)
+    mask_windows = mask_windows.squeeze(-1)
+    attention_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+    attention_mask = attention_mask.masked_fill(attention_mask != 0, float(-100.0)).masked_fill(
+        attention_mask == 0, float(0.0)
+    )
+    return attention_mask
+
+
+class OmnivoreSwinTransformerStage(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        num_heads,
+        window_size=(1, 7, 7),
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        dropout_rate=0.0,
+        attention_dropout_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = tuple(i // 2 for i in window_size)
+        self.depth = depth
+
+        # build layers
+        self.layers = nn.ModuleList(
+            [
+                OmnivoreSwinTransformer3DLayer(
+                    dim=dim,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    shift_size=(0, 0, 0) if (i % 2 == 0) else self.shift_size,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    dropout_rate=dropout_rate,
+                    attention_dropout_rate=attention_dropout_rate,
+                    drop_path_rate=drop_path_rate[i] if isinstance(drop_path_rate, list) else drop_path_rate,
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.downsample = downsample
+        if self.downsample is not None:
+            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+
+    def forward(self, hidden_state, use_checkpoint=False, height=None, width=None, use_seg=False):
+        if use_seg:
+            return self.forward_seg(hidden_state, height, width)
+        batch_size, channels, D, height, width = hidden_state.shape
+        window_size, shift_size = get_window_size((D, height, width), self.window_size, self.shift_size)
+        hidden_state = hidden_state.permute(0, 2, 3, 4, 1)
+
+        Dp = int(np.ceil(D / window_size[0])) * window_size[0]
+        Hp = int(np.ceil(height / window_size[1])) * window_size[1]
+        Wp = int(np.ceil(width / window_size[2])) * window_size[2]
+
+        attention_mask = compute_mask(Dp, Hp, Wp, window_size, shift_size, hidden_state.device)
+
+        for layer in self.layers:
+            hidden_state = layer(hidden_state, attention_mask, use_checkpoint=use_checkpoint)
+        hidden_state = hidden_state.view(batch_size, D, height, width, -1)
+
+        if self.downsample is not None:
+            hidden_state = self.downsample(hidden_state)
+
+        hidden_state = hidden_state.permute(0, 4, 1, 2, 3)
+
+        return hidden_state
+
+    def forward_seg(self, hidden_state, height, width):
+
+        Hp = int(np.ceil(height / self.window_size[1])) * self.window_size[1]
+        Wp = int(np.ceil(width / self.window_size[2])) * self.window_size[2]
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=hidden_state.device)  # 1 Hp Wp 1
+        h_slices = (
+            slice(0, -self.window_size[1]),
+            slice(-self.window_size[1], -self.shift_size[1]),
+            slice(-self.shift_size[1], None),
+        )
+        w_slices = (
+            slice(0, -self.window_size[2]),
+            slice(-self.window_size[2], -self.shift_size[2]),
+            slice(-self.shift_size[2], None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition_image(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size[1] * self.window_size[2])
+        attention_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attention_mask = attention_mask.masked_fill(attention_mask != 0, float(-100.0)).masked_fill(
+            attention_mask == 0, float(0.0)
+        )
+
+        for layer in self.layers:
+            layer.height, layer.width = height, width
+            if hidden_state.ndim == 4:
+                batch_size, D, channels, seq_len = hidden_state.shape
+                assert seq_len == height * width, "input feature has wrong size"
+                hidden_state = hidden_state.reshape(batch_size, D, channels, height, width)
+                hidden_state = hidden_state.permute(0, 1, 3, 4, 2)
+            assert hidden_state.shape[2] == height
+            assert hidden_state.shape[3] == width
+            hidden_state = layer(hidden_state, attention_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(hidden_state, height, width)
+            Wh, Ww = (height + 1) // 2, (width + 1) // 2
+            return hidden_state, height, width, x_down, Wh, Ww
+        else:
+            return hidden_state, height, width, hidden_state, height, width
+
+
+class OmnivorePatchEmbeddings3D(nn.Module):
+    """Video to Patch Embedding"""
+
+    def __init__(
+        self,
+        patch_size=(2, 4, 4),
+        input_channels=3,
+        embed_dim=96,
+        norm_layer=None,
+        additional_variable_channels=None,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.input_channels = input_channels
+        self.embed_dim = embed_dim
+        self.additional_variable_channels = additional_variable_channels
+
+        self.projection = nn.Conv3d(input_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if additional_variable_channels:
+            # we create var_proj separately from proj
+            # this makes it convenient to ignore var_proj on downstream tasks
+            # where we only use RGB
+            self.var_projection = [
+                nn.Conv3d(x, embed_dim, kernel_size=patch_size, stride=patch_size)
+                for x in additional_variable_channels
+            ]
+            self.var_projection = nn.ModuleList(self.var_projection)
+
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def run_variable_channel_forward(self, hidden_state):
+        sidx = 0
+        out = None
+        for idx in range(len(self.additional_variable_channels)):
+            eidx = sidx + self.additional_variable_channels[idx]
+            c_out = self.var_projection[idx](hidden_state[:, sidx:eidx, ...])
+            if idx == 0:
+                out = c_out
+            else:
+                out += c_out
+            sidx = eidx
+        return out
+
+    def forward(self, hidden_state):
+        _, _, D, height, width = hidden_state.size()
+        if width % self.patch_size[2] != 0:
+            hidden_state = F.pad(hidden_state, (0, self.patch_size[2] - width % self.patch_size[2]))
+        if height % self.patch_size[1] != 0:
+            hidden_state = F.pad(hidden_state, (0, 0, 0, self.patch_size[1] - height % self.patch_size[1]))
+        if D % self.patch_size[0] != 0:
+            hidden_state = F.pad(hidden_state, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
+
+        if self.additional_variable_channels:
+            hidden_state_rgb = hidden_state[:, :3, ...]
+            hidden_state_rem = hidden_state[:, 3:, ...]
+            hidden_state_rgb = self.projection(hidden_state_rgb)
+            if hidden_state.shape[1] > 3:
+                hidden_state_rem = self.run_variable_channel_forward(hidden_state_rem)
+                hidden_state = hidden_state_rgb + hidden_state_rem
+            else:
+                hidden_state = hidden_state_rgb
+        else:
+            hidden_state = self.projection(hidden_state)  # B C D Wh Ww
+        if self.norm is not None:
+            D, Wh, Ww = hidden_state.size(2), hidden_state.size(3), hidden_state.size(4)
+            hidden_state = hidden_state.flatten(2).transpose(1, 2)
+            hidden_state = self.norm(hidden_state)
+            hidden_state = hidden_state.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
+
+        return hidden_state
+
+
+class OmnivoreSwinTransformer3DModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.im2vid = OmnivoreIm2Video()
+        self.num_stages = len(self.config.depths)
+        self.patch_size = self.config.patch_size
+        self.input_channels = self.config.input_channels
+        self.embed_dim = self.config.embed_dim
+        self.depths = self.config.depths
+        self.num_heads = self.config.num_heads
+        self.window_size = self.config.window_size
+        self.mlp_ratio = self.config.mlp_ratio
+        self.qkv_bias = self.config.qkv_bias
+        self.qk_scale = self.config.qk_scale
+        self.dropout_rate = self.config.dropout_rate
+        self.attention_dropout_rate = self.config.attention_dropout_rate
+        self.drop_path_rate = self.config.drop_path_rate
+        self.norm_layer = nn.LayerNorm
+        self.patch_norm = self.config.patch_norm
+        self.frozen_stages = self.config.frozen_stages
+        self.depth_patch_embed_separate_params = True
+        self.depth_mode = self.config.depth_mode
+        depth_chans = None
+        assert self.input_channels == 3, "Only 3 channels supported"
+
+        # split image into non-overlapping patches
+        self.patch_embed = OmnivorePatchEmbeddings3D(
+            patch_size=self.patch_size,
+            input_channels=self.input_channels,
+            embed_dim=self.embed_dim,
+            norm_layer=self.norm_layer if self.patch_norm else None,
+        )
+
+        if self.depth_mode is not None:
+            msg = f"Using depth mode {self.depth_mode}"
+            logger.info(msg)
+            assert self.depth_mode in ["separate_d_tokens", "summed_rgb_d_tokens", "rgbd"]
+            if self.depth_mode in ["separate_d_tokens", "summed_rgb_d_tokens"]:
+                depth_chans = 1
+                assert self.depth_patch_embed_separate_params, "separate tokenization needs separate parameters"
+                if self.depth_mode == "separate_d_tokens":
+                    raise NotImplementedError()
+            else:
+                assert self.depth_mode == "rgbd"
+                depth_chans = 4
+
+            self.depth_patch_embed_separate_params = self.depth_patch_embed_separate_params
+
+            if self.depth_patch_embed_separate_params:
+                self.depth_patch_embed = OmnivorePatchEmbeddings3D(
+                    patch_size=self.patch_size,
+                    input_channels=depth_chans,
+                    embed_dim=self.embed_dim,
+                    norm_layer=self.norm_layer if self.patch_norm else None,
+                )
+            else:
+                del self.patch_embed
+                assert depth_chans == 4
+                logger.info("Certain channels of patch projection may not be used in forward pass")
+                logger.info("Make sure config.DISTRIBUTED.FIND_UNUSED_PARAMETERS is set to True")
+                self.patch_embed = OmnivorePatchEmbeddings3D(
+                    patch_size=self.patch_size,
+                    input_channels=3,
+                    embed_dim=self.embed_dim,
+                    additional_variable_channels=[1],
+                    norm_layer=self.norm_layer if self.patch_norm else None,
+                )
+
+        self.pos_drop = nn.Dropout(p=self.dropout_rate)
+
+        # stochastic depth
+        dpr = [
+            x.item() for x in torch.linspace(0, self.drop_path_rate, sum(self.depths))
+        ]  # stochastic depth decay rule
+
+        # build stages
+        self.stages = nn.ModuleList()
+        for stage in range(self.num_stages):
+            stage_module = OmnivoreSwinTransformerStage(
+                dim=int(self.embed_dim * 2**stage),
+                depth=self.depths[stage],
+                num_heads=self.num_heads[stage],
+                window_size=self.window_size,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=self.qkv_bias,
+                qk_scale=self.qk_scale,
+                dropout_rate=self.dropout_rate,
+                attention_dropout_rate=self.attention_dropout_rate,
+                drop_path_rate=dpr[sum(self.depths[:stage]) : sum(self.depths[: stage + 1])],
+                norm_layer=self.norm_layer,
+                downsample=OmnivorePatchMerging if stage < self.num_stages - 1 else None,
+            )
+            self.stages.append(stage_module)
+
+        self.num_features = int(self.embed_dim * 2 ** (self.num_stages - 1))
+        self.norm = self.norm_layer(self.num_features)
+        self._freeze_stages()
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            self.patch_embed.eval()
+            for param in self.patch_embed.parameters():
+                param.requires_grad = False
+
+        if self.frozen_stages >= 1:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+    def _apply_norm(self, x):
+        x = x.permute(0, 2, 3, 4, 1)
+        x = self.norm(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        return x
+
+    def forward_intermediate_features(self, stage_outputs, out_feat_keys):
+        """
+        Inputs
+        - stage_outputs: list of features without self.norm() applied to them
+        - out_feat_keys: list of feature names (str)
+                         specified as "stage<int>" for feature with norm or "interim<int>" for feature without norm
+        """
+        out_features = []
+        for key in out_feat_keys:
+            if key.startswith("stage"):
+                rep = "stage"
+            elif key.startswith("interim"):
+                rep = "interim"
+            else:
+                raise ValueError(f"Invalid key {key}")
+            idx = int(key.replace(rep, ""))
+            feat = stage_outputs[idx]
+            if rep == "stage":
+                feat = self._apply_norm(feat)
+            out_features.append(feat)
+        return out_features
+
+    def get_patch_embedding(self, hidden_state):
+        assert hidden_state.ndim == 5
+        has_depth = hidden_state.shape[1] == 4
+
+        if has_depth:
+            if self.depth_mode in ["summed_rgb_d_tokens"]:
+                hidden_state_rgb = hidden_state[:, :3, ...]
+                hidden_state_d = hidden_state[:, 3:, ...]
+                hidden_state_d = self.depth_patch_embed(hidden_state_d)
+                hidden_state_rgb = self.patch_embed(hidden_state_rgb)
+                # sum the two sets of tokens
+                hidden_state = hidden_state_rgb + hidden_state_d
+            elif self.depth_mode == "rgbd":
+                if self.depth_patch_embed_separate_params:
+                    hidden_state = self.depth_patch_embed(hidden_state)
+                else:
+                    hidden_state = self.patch_embed(hidden_state)
+            else:
+                logger.info("Depth mode %s not supported" % self.depth_mode)
+                raise NotImplementedError()
+        else:
+            hidden_state = self.patch_embed(hidden_state)
+        return hidden_state
+
+    def forward(
+        self, hidden_state, out_feat_keys=None, use_checkpoint=False, output_hidden_states=False, return_dict=True
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        hidden_state = self.im2vid(hidden_state)
+        hidden_state = self.get_patch_embedding(hidden_state)
+        hidden_state = self.pos_drop(hidden_state)
+
+        stage_outputs = []
+
+        for stage in self.stages:
+            hidden_state = stage(hidden_state.contiguous(), use_checkpoint=use_checkpoint)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_state,)
+            stage_outputs.append(hidden_state)
+
+        if out_feat_keys is not None and len(out_feat_keys) > 0:
+            final_hidden_state = self.forward_intermediate_features(stage_outputs, out_feat_keys)
+        else:
+            hidden_state = self._apply_norm(hidden_state)
+            # Mean over the spatiotemporal dimensions
+            hidden_state = torch.mean(hidden_state, [-3, -2, -1])
+
+            final_hidden_state = hidden_state
+
+        if not return_dict:
+            return tuple(v for v in [final_hidden_state, all_hidden_states] if v is not None)
+        return BaseModelOutputWithNoAttention(last_hidden_state=final_hidden_state, hidden_states=all_hidden_states)
+
+    def train(self, mode=True):
+        """Convert the model into training mode while keep layers freezed."""
+        super(OmnivoreSwinTransformer3DModel, self).train(mode)
+        self._freeze_stages()
+
+
+class OmnivoreImageClassificationHead(nn.Module):
+    def __init__(self, in_features=1024, out_features=1000, bias=True):
+        super().__init__()
+        self.image_head = nn.Linear(in_features, out_features, bias)
+
+    def forward(self, hidden_state):
+        logits = self.image_head(hidden_state)
+        return logits
+
+
+class OmnivoreVideoClassificationHead(nn.Module):
+    def __init__(self, in_features=1024, out_features=400, bias=True):
+        super().__init__()
+        self.video_head = nn.Linear(in_features, out_features, bias)
+        self.dropout = nn.Dropout(p=0.5)
+
+    def forward(self, hidden_state):
+        logits = self.video_head(hidden_state)
+        logits = self.dropout(logits)
+        return logits
+
+
+class OmnivoreRGBDClassificationHead(nn.Module):
+    def __init__(self, in_features=1024, out_features=19, bias=True):
+        super().__init__()
+        self.rgbd_head = nn.Linear(in_features, out_features, bias)
+
+    def forward(self, hidden_state):
+        logits = self.rgbd_head(hidden_state)
+        return logits
+
+
+class OmnivorePreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = OmnivoreConfig
+    base_model_prefix = "omnivore"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, OmnivoreModel):
+            module.gradient_checkpointing = value
+
+
+OMNIVORE_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
+    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+
+    Parameters:
+        config ([`OmnivoreConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+OMNIVORE_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoFeatureExtractor`]. See
+            [`AutoFeatureExtractor.__call__`] for details.
+
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Omnivore model outputting raw features without any specific head on top.",
+    OMNIVORE_START_DOCSTRING,
+)
+class OmnivoreModel(OmnivorePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = OmnivoreSwinTransformer3DModel(config)
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(OMNIVORE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        processor_class=_FEAT_EXTRACTOR_FOR_DOC,
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPoolingAndNoAttention,
+        config_class=_CONFIG_FOR_DOC,
+        modality="vision",
+        expected_output=_EXPECTED_OUTPUT_SHAPE,
+    )
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        outputs = self.model(pixel_values)
+        last_hidden_state = outputs[0]
+        # global average pooling, (N, C, D, H, W) -> (N, C)
+        pooled_output = last_hidden_state.mean([-1])
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + outputs[1:]
+
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=outputs.hidden_states,
+        )
+
+
+@add_start_docstrings(
+    """
+    Omnivore Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+    ImageNet.
+    """,
+    OMNIVORE_START_DOCSTRING,
+)
+class OmnivoreForImageClassification(OmnivorePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.num_image_labels = config.num_image_labels or config.num_labels
+        self.num_video_labels = config.num_video_labels or config.num_labels
+        self.num_rgbd_labels = config.num_rgbd_labels or config.num_labels
+        self.omnivore = OmnivoreModel(config)
+        self.image_classifier = OmnivoreImageClassificationHead(config.head_dim_in, self.num_image_labels)
+        self.rgbd_classifier = OmnivoreRGBDClassificationHead(config.head_dim_in, self.num_rgbd_labels)
+        self.video_classifier = OmnivoreVideoClassificationHead(config.head_dim_in, self.num_video_labels)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(OMNIVORE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        processor_class=_FEAT_EXTRACTOR_FOR_DOC,
+        checkpoint=_IMAGE_CLASS_CHECKPOINT,
+        output_type=ImageClassifierOutputWithNoAttention,
+        config_class=_CONFIG_FOR_DOC,
+        modality="vision",
+        expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
+    )
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        pixel_input_type: str = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        r"""
+        pixel_input_type (`str`):
+            Which classification head to use for the classification of given pixel_values
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        Returns:
+
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.omnivore(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
+        sequence_output = outputs[0]
+
+        logits = None
+        if pixel_input_type == "image":
+            logits = self.image_classifier(sequence_output)
+
+        if pixel_input_type == "video":
+            logits = self.video_classifier(sequence_output)
+
+        if pixel_input_type == "rgbd":
+            logits = self.rgbd_classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+        )