codingutf8 2022 meta platforms inc and the huggingface inc team 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 orglicenseslicense2 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 maskformer swin transformer the reason swin transformer is implemented here is because maskformer uses the hidden states before downsampling which is different from the default swin transformer import collections abc import math from dataclasses import dataclass from typing import optional tuple import torch from torch import tensor nn from activations import act2fn from fileutils import modeloutput from modelingoutputs import backboneoutput from modelingutils import pretrainedmodel from pytorchutils import findpruneableheadsandindices meshgrid prunelinearlayer from utils backboneutils import backbonemixin from configurationmaskformerswin import maskformerswinconfig dataclass class maskformerswinmodeloutputwithpoolingmodeloutput lasthiddenstate torch floattensor none pooleroutput torch floattensor none hiddenstates optionaltupletorch floattensor none hiddenstatesspatialdimensions tupletupleint int none attentions optionaltupletorch floattensor none dataclass class maskformerswinbasemodeloutputmodeloutput lasthiddenstate torch floattensor none hiddenstates optionaltupletorch floattensor none hiddenstatesspatialdimensions tupletupleint int none attentions optionaltupletorch floattensor none copied from transformers models swin modelingswin windowpartition def windowpartitioninputfeature windowsize batchsize height width numchannels inputfeature shape inputfeature inputfeature view batchsize height windowsize windowsize width windowsize windowsize numchannels windows inputfeature permute0 1 3 2 4 5 contiguous view1 windowsize windowsize numchannels return windows copied from transformers models swin modelingswin windowreverse def windowreversewindows windowsize height width numchannels windows shape1 windows windows view1 height windowsize width windowsize windowsize windowsize numchannels windows windows permute0 1 3 2 4 5 contiguous view1 height width numchannels return windows copied from transformers models swin modelingswin droppath def droppathinput torch tensor dropprob float 0 0 training bool false torch tensor if dropprob 0 0 or not training return input keepprob 1 dropprob shape input shape0 1 input ndim 1 work with diff dim tensors not just 2d convnets randomtensor keepprob torch randshape dtypeinput dtype deviceinput device randomtensor floor binarize output input divkeepprob randomtensor return output class maskformerswinembeddingsnn module def initself config super init self patchembeddings maskformerswinpatchembeddingsconfig numpatches self patchembeddings numpatches self patchgrid self patchembeddings gridsize if config useabsoluteembeddings self positionembeddings nn parametertorch zeros1 numpatches 1 config embeddim else self positionembeddings none self norm nn layernormconfig embeddim self dropout nn dropoutconfig hiddendropoutprob def forwardself pixelvalues embeddings outputdimensions self patchembeddingspixelvalues embeddings self normembeddings if self positionembeddings is not none embeddings embeddings self positionembeddings embeddings self dropoutembeddings return embeddings outputdimensions copied from transformers models swin modelingswin swinpatchembeddings class maskformerswinpatchembeddingsnn module def initself config super init imagesize patchsize config imagesize config patchsize numchannels hiddensize config numchannels config embeddim imagesize imagesize if isinstanceimagesize collections abc iterable else imagesize imagesize patchsize patchsize if isinstancepatchsize collections abc iterable else patchsize patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 self imagesize imagesize self patchsize patchsize self numchannels numchannels self numpatches numpatches self gridsize imagesize0 patchsize0 imagesize1 patchsize1 self projection nn conv2dnumchannels hiddensize kernelsizepatchsize stridepatchsize def maybepadself pixelvalues height width if width self patchsize1 0 padvalues 0 self patchsize1 width self patchsize1 pixelvalues nn functional padpixelvalues padvalues if height self patchsize0 0 padvalues 0 0 0 self patchsize0 height self patchsize0 pixelvalues nn functional padpixelvalues padvalues return pixelvalues def forwardself pixelvalues optionaltorch floattensor tupletorch tensor tupleint numchannels height width pixelvalues shape if numchannels self numchannels raise valueerror make sure that the channel dimension of the pixel values match with the one set in the configuration pad the input to be divisible by self patchsize if needed pixelvalues self maybepadpixelvalues height width embeddings self projectionpixelvalues height width embeddings shape outputdimensions height width embeddings embeddings flatten2 transpose1 2 return embeddings outputdimensions copied from transformers models swin modelingswin swinpatchmerging class maskformerswinpatchmergingnn module def initself inputresolution tupleint dim int normlayer nn module nn layernorm none super init self inputresolution inputresolution self dim dim self reduction nn linear4 dim 2 dim biasfalse self norm normlayer4 dim def maybepadself inputfeature height width shouldpad height 2 1 or width 2 1 if shouldpad padvalues 0 0 0 width 2 0 height 2 inputfeature nn functional padinputfeature padvalues return inputfeature def forwardself inputfeature torch tensor inputdimensions tupleint int torch tensor height width inputdimensions dim is height width batchsize dim numchannels inputfeature shape inputfeature inputfeature viewbatchsize height width numchannels pad input to be disible by width and height if needed inputfeature self maybepadinputfeature height width batchsize height2 width2 numchannels inputfeature0 inputfeature 0 2 0 2 batchsize height2 width2 numchannels inputfeature1 inputfeature 1 2 0 2 batchsize height2 width2 numchannels inputfeature2 inputfeature 0 2 1 2 batchsize height2 width2 numchannels inputfeature3 inputfeature 1 2 1 2 batchsize height2 width2 4numchannels inputfeature torch catinputfeature0 inputfeature1 inputfeature2 inputfeature3 1 inputfeature inputfeature viewbatchsize 1 4 numchannels batchsize height2width2 4c inputfeature self norminputfeature inputfeature self reductioninputfeature return inputfeature copied from transformers models swin modelingswin swindroppath with swinmaskformerswin class maskformerswindroppathnn module copied from transformers models swin modelingswin swinselfattention with swinmaskformerswin get pairwise relative position index for each token inside the window take the dot product between query and key to get the raw attention scores apply the attention mask is precomputed for all layers in maskformerswinmodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models swin modelingswin swinselfoutput with swinmaskformerswin copied from transformers models swin modelingswin swinattention with swinmaskformerswin prune linear layers update hyper params and store pruned heads copied from transformers models swin modelingswin swinintermediate with swinmaskformerswin copied from transformers models swin modelingswin swinoutput with swinmaskformerswin calculate attention mask for swmsa pad hiddenstates to multiples of window size cyclic shift partition windows reverse cyclic shift copied from transformers models swin modelingswin swinstage init with swinmaskformerswin patch merging layer copied from transformers models swin modelingswin swinencoder init with swinmaskformerswin copied from transformers models swin modelingswin swinpretrainedmodel with swinmaskformerswin swinmodel an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights if isinstancemodule nn linear nn conv2d slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 module weight data normalmean0 0 stdself config initializerrange if module bias is not none module bias data zero elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 class maskformerswinmodelmaskformerswinpretrainedmodel def initself config addpoolinglayertrue super initconfig self config config self numlayers lenconfig depths self numfeatures intconfig embeddim 2 self numlayers 1 self embeddings maskformerswinembeddingsconfig self encoder maskformerswinencoderconfig self embeddings patchgrid self layernorm nn layernormself numfeatures epsconfig layernormeps self pooler nn adaptiveavgpool1d1 if addpoolinglayer else none def getinputembeddingsself return self embeddings patchembeddings def pruneheadsself headstoprune for layer heads in headstoprune items self encoder layerlayer attention pruneheadsheads def forward self pixelvaluesnone headmasknone outputattentionsnone outputhiddenstatesnone returndictnone outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict if pixelvalues is none raise valueerroryou have to specify pixelvalues prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x nheads x n x n input headmask has shape numheads or numhiddenlayers x numheads and headmask is converted to shape numhiddenlayers x batch x numheads x seqlength x seqlength headmask self getheadmaskheadmask lenself config depths embeddingoutput inputdimensions self embeddingspixelvalues encoderoutputs self encoder embeddingoutput inputdimensions headmaskheadmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict sequenceoutput encoderoutputs lasthiddenstate if returndict else encoderoutputs0 sequenceoutput self layernormsequenceoutput pooledoutput none if self pooler is not none pooledoutput self poolersequenceoutput transpose1 2 pooledoutput torch flattenpooledoutput 1 if not returndict return sequenceoutput pooledoutput encoderoutputs1 hiddenstatesspatialdimensions inputdimensions encoderoutputs hiddenstatesspatialdimensions return maskformerswinmodeloutputwithpooling lasthiddenstatesequenceoutput pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates hiddenstatesspatialdimensionshiddenstatesspatialdimensions attentionsencoderoutputs attentions class maskformerswinbackbonemaskformerswinpretrainedmodel backbonemixin def initself config maskformerswinconfig super initconfig super initbackboneconfig self model maskformerswinmodelconfig if stem in self outfeatures raise valueerrorthis backbone does not support stem in the outfeatures self numfeatures config embeddim intconfig embeddim 2i for i in rangelenconfig depths self hiddenstatesnorms nn modulelist nn layernormnumchannels for numchannels in self numfeatures1 initialize weights and apply final processing self postinit def forward self pixelvalues tensor outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none backboneoutput returndict returndict if returndict is not none else self config usereturndict outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates outputattentions outputattentions if outputattentions is not none else self config outputattentions outputs self model pixelvalues outputhiddenstatestrue outputattentionsoutputattentions returndicttrue we skip the stem hiddenstates outputs hiddenstates1 we need to reshape the hidden states to their original spatial dimensions spatial dimensions contains all the heights and widths of each stage including after the embeddings spatialdimensions tupletupleint int outputs hiddenstatesspatialdimensions featuremaps for i hiddenstate stage height width in enumerate ziphiddenstates self stagenames1 spatialdimensions norm self hiddenstatesnormsi the last element corespond to the layer s last block output but before patch merging hiddenstateunpolled hiddenstate1 hiddenstatenorm normhiddenstateunpolled the pixel decoder fpn expects 3d tensors features batchsize hiddensize hiddenstatenorm shape reshape b h w d b d h w hiddenstatepermuted hiddenstatenorm permute0 2 1 viewbatchsize hiddensize height width contiguous if stage in self outfeatures featuremaps hiddenstatepermuted if not returndict output featuremaps if outputhiddenstates output outputs hiddenstates if outputattentions output outputs attentions return output return backboneoutput featuremapsfeaturemaps hiddenstatesoutputs hiddenstates if outputhiddenstates else none attentionsoutputs attentions coding utf 8 2022 meta platforms inc and the huggingface inc team 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 maskformer swin transformer the reason swin transformer is implemented here is because maskformer uses the hidden states before downsampling which is different from the default swin transformer class for maskformerswinmodel s outputs that also contains the spatial dimensions of the hidden states args last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model pooler_output torch floattensor of shape batch_size hidden_size last layer hidden state after a mean pooling operation hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs hidden_states_spatial_dimensions tuple tuple int int optional a tuple containing the spatial dimension of each hidden_state needed to reshape the hidden_states to batch channels height width due to padding their spatial size cannot be inferred before the forward method attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads class for swinencoder s outputs args last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs hidden_states_spatial_dimensions tuple tuple int int optional a tuple containing the spatial dimension of each hidden_state needed to reshape the hidden_states to batch channels height width due to padding their spatial size cannot inferred before the forward method attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads copied from transformers models swin modeling_swin window_partition partitions the given input into windows copied from transformers models swin modeling_swin window_reverse merges windows to produce higher resolution features copied from transformers models swin modeling_swin drop_path drop paths stochastic depth per sample when applied in main path of residual blocks comment by ross wightman 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 work with diff dim tensors not just 2d convnets binarize construct the patch and position embeddings copied from transformers models swin modeling_swin swinpatchembeddings this class turns pixel_values of shape batch_size num_channels height width into the initial hidden_states patch embeddings of shape batch_size seq_length hidden_size to be consumed by a transformer pad the input to be divisible by self patch_size if needed copied from transformers models swin modeling_swin swinpatchmerging patch merging layer args input_resolution tuple int resolution of input feature dim int number of input channels norm_layer nn module optional defaults to nn layernorm normalization layer class dim is height width pad input to be disible by width and height if needed batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 num_channels batch_size height 2 width 2 4 num_channels batch_size height 2 width 2 4 c copied from transformers models swin modeling_swin swindroppath with swin maskformerswin drop paths stochastic depth per sample when applied in main path of residual blocks copied from transformers models swin modeling_swin swinselfattention with swin maskformerswin get pair wise relative position index for each token inside the window take the dot product between query and key to get the raw attention scores apply the attention mask is precomputed for all layers in maskformerswinmodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models swin modeling_swin swinselfoutput with swin maskformerswin copied from transformers models swin modeling_swin swinattention with swin maskformerswin prune linear layers update hyper params and store pruned heads add attentions if we output them copied from transformers models swin modeling_swin swinintermediate with swin maskformerswin copied from transformers models swin modeling_swin swinoutput with swin maskformerswin calculate attention mask for sw msa pad hidden_states to multiples of window size cyclic shift partition windows add self attentions if we output attention weights b height width c reverse cyclic shift copied from transformers models swin modeling_swin swinstage __init__ with swin maskformerswin patch merging layer copied from transformers models swin modeling_swin swinencoder __init__ with swin maskformerswin copied from transformers models swin modeling_swin swinpretrainedmodel with swin maskformerswin swin model an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x n x n input head_mask has shape num_heads or num_hidden_layers x num_heads and head_mask is converted to shape num_hidden_layers x batch x num_heads x seq_length x seq_length maskformerswin backbone designed especially for the maskformer framework this classes reshapes hidden_states from batch_size sequence_length hidden_size to batch_size num_channels height width it also adds additional layernorms after each stage args config maskformerswinconfig the configuration used by maskformerswinmodel initialize weights and apply final processing we skip the stem we need to reshape the hidden states to their original spatial dimensions spatial dimensions contains all the heights and widths of each stage including after the embeddings the last element corespond to the layer s last block output but before patch merging the pixel decoder fpn expects 3d tensors features reshape b h w d b d h w
import collections.abc
import math
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import Tensor, nn

from ...activations import ACT2FN
from ...file_utils import ModelOutput
from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
from ...utils.backbone_utils import BackboneMixin
from .configuration_maskformer_swin import MaskFormerSwinConfig


@dataclass
class MaskFormerSwinModelOutputWithPooling(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    pooler_output: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states_spatial_dimensions: Tuple[Tuple[int, int]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class MaskFormerSwinBaseModelOutput(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states_spatial_dimensions: Tuple[Tuple[int, int]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None



def window_partition(input_feature, window_size):
    
    batch_size, height, width, num_channels = input_feature.shape
    input_feature = input_feature.view(
        batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
    )
    windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
    return windows



def window_reverse(windows, window_size, height, width):
    
    num_channels = windows.shape[-1]
    windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)
    return windows



def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
    
    if drop_prob == 0.0 or not training:
        return input
    keep_prob = 1 - drop_prob
    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  
    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
    random_tensor.floor_()  
    output = input.div(keep_prob) * random_tensor
    return output


class MaskFormerSwinEmbeddings(nn.Module):
    

    def __init__(self, config):
        super().__init__()

        self.patch_embeddings = MaskFormerSwinPatchEmbeddings(config)
        num_patches = self.patch_embeddings.num_patches
        self.patch_grid = self.patch_embeddings.grid_size

        if config.use_absolute_embeddings:
            self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim))
        else:
            self.position_embeddings = None

        self.norm = nn.LayerNorm(config.embed_dim)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, pixel_values):
        embeddings, output_dimensions = self.patch_embeddings(pixel_values)
        embeddings = self.norm(embeddings)

        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings)

        return embeddings, output_dimensions



class MaskFormerSwinPatchEmbeddings(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        image_size, patch_size = config.image_size, config.patch_size
        num_channels, hidden_size = config.num_channels, config.embed_dim
        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.num_patches = num_patches
        self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])

        self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)

    def maybe_pad(self, pixel_values, height, width):
        if width % self.patch_size[1] != 0:
            pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
        if height % self.patch_size[0] != 0:
            pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
        return pixel_values

    def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
        _, num_channels, height, width = pixel_values.shape
        if num_channels != self.num_channels:
            raise ValueError(
                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
            )
        
        pixel_values = self.maybe_pad(pixel_values, height, width)
        embeddings = self.projection(pixel_values)
        _, _, height, width = embeddings.shape
        output_dimensions = (height, width)
        embeddings = embeddings.flatten(2).transpose(1, 2)

        return embeddings, output_dimensions



class MaskFormerSwinPatchMerging(nn.Module):
    

    def __init__(self, input_resolution: Tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def maybe_pad(self, input_feature, height, width):
        should_pad = (height % 2 == 1) or (width % 2 == 1)
        if should_pad:
            pad_values = (0, 0, 0, width % 2, 0, height % 2)
            input_feature = nn.functional.pad(input_feature, pad_values)

        return input_feature

    def forward(self, input_feature: torch.Tensor, input_dimensions: Tuple[int, int]) -> torch.Tensor:
        height, width = input_dimensions
        
        batch_size, dim, num_channels = input_feature.shape

        input_feature = input_feature.view(batch_size, height, width, num_channels)
        
        input_feature = self.maybe_pad(input_feature, height, width)
        
        input_feature_0 = input_feature[:, 0::2, 0::2, :]
        
        input_feature_1 = input_feature[:, 1::2, 0::2, :]
        
        input_feature_2 = input_feature[:, 0::2, 1::2, :]
        
        input_feature_3 = input_feature[:, 1::2, 1::2, :]
        
        input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
        input_feature = input_feature.view(batch_size, -1, 4 * num_channels)  

        input_feature = self.norm(input_feature)
        input_feature = self.reduction(input_feature)

        return input_feature



class MaskFormerSwinDropPath(nn.Module):
    

    def __init__(self, drop_prob: Optional[float] = None) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return drop_path(hidden_states, self.drop_prob, self.training)

    def extra_repr(self) -> str:
        return "p={}".format(self.drop_prob)



class MaskFormerSwinSelfAttention(nn.Module):
    def __init__(self, config, dim, num_heads, window_size):
        super().__init__()
        if dim % num_heads != 0:
            raise ValueError(
                f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
            )

        self.num_attention_heads = num_heads
        self.attention_head_size = int(dim / num_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.window_size = (
            window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
        )

        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads)
        )

        
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))
        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
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)
        self.register_buffer("relative_position_index", relative_position_index)

        self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
        self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
        self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        batch_size, dim, num_channels = hidden_states.shape
        mixed_query_layer = self.query(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        query_layer = self.transpose_for_scores(mixed_query_layer)

        
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)]
        relative_position_bias = relative_position_bias.view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
        )

        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
        attention_scores = attention_scores + relative_position_bias.unsqueeze(0)

        if attention_mask is not None:
            
            mask_shape = attention_mask.shape[0]
            attention_scores = attention_scores.view(
                batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim
            )
            attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)
            attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)

        
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        
        
        attention_probs = self.dropout(attention_probs)

        
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs



class MaskFormerSwinSelfOutput(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)

        return hidden_states



class MaskFormerSwinAttention(nn.Module):
    def __init__(self, config, dim, num_heads, window_size):
        super().__init__()
        self.self = MaskFormerSwinSelfAttention(config, dim, num_heads, window_size)
        self.output = MaskFormerSwinSelfOutput(config, dim)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        self_outputs = self.self(hidden_states, attention_mask, head_mask, output_attentions)
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  
        return outputs



class MaskFormerSwinIntermediate(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states



class MaskFormerSwinOutput(nn.Module):
    def __init__(self, config, dim):
        super().__init__()
        self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


class MaskFormerSwinLayer(nn.Module):
    def __init__(self, config, dim, input_resolution, num_heads, shift_size=0):
        super().__init__()
        self.shift_size = shift_size
        self.window_size = config.window_size
        self.input_resolution = input_resolution
        self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
        self.attention = MaskFormerSwinAttention(config, dim, num_heads, self.window_size)
        self.drop_path = (
            MaskFormerSwinDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
        )
        self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
        self.intermediate = MaskFormerSwinIntermediate(config, dim)
        self.output = MaskFormerSwinOutput(config, dim)

    def get_attn_mask(self, input_resolution):
        if self.shift_size > 0:
            
            height, width = input_resolution
            img_mask = torch.zeros((1, height, width, 1))
            height_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            width_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            count = 0
            for height_slice in height_slices:
                for width_slice in width_slices:
                    img_mask[:, height_slice, width_slice, :] = count
                    count += 1

            mask_windows = window_partition(img_mask, self.window_size)
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None
        return attn_mask

    def maybe_pad(self, hidden_states, height, width):
        pad_left = pad_top = 0
        pad_rigth = (self.window_size - width % self.window_size) % self.window_size
        pad_bottom = (self.window_size - height % self.window_size) % self.window_size
        pad_values = (0, 0, pad_left, pad_rigth, pad_top, pad_bottom)
        hidden_states = nn.functional.pad(hidden_states, pad_values)
        return hidden_states, pad_values

    def forward(self, hidden_states, input_dimensions, head_mask=None, output_attentions=False):
        height, width = input_dimensions
        batch_size, dim, channels = hidden_states.size()
        shortcut = hidden_states

        hidden_states = self.layernorm_before(hidden_states)
        hidden_states = hidden_states.view(batch_size, height, width, channels)
        
        hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)

        _, height_pad, width_pad, _ = hidden_states.shape
        
        if self.shift_size > 0:
            shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_hidden_states = hidden_states

        
        hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
        hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
        attn_mask = self.get_attn_mask((height_pad, width_pad))
        if attn_mask is not None:
            attn_mask = attn_mask.to(hidden_states_windows.device)

        self_attention_outputs = self.attention(
            hidden_states_windows, attn_mask, head_mask, output_attentions=output_attentions
        )

        attention_output = self_attention_outputs[0]

        outputs = self_attention_outputs[1:]  

        attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
        shifted_windows = window_reverse(
            attention_windows, self.window_size, height_pad, width_pad
        )  

        
        if self.shift_size > 0:
            attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            attention_windows = shifted_windows

        was_padded = pad_values[3] > 0 or pad_values[5] > 0
        if was_padded:
            attention_windows = attention_windows[:, :height, :width, :].contiguous()

        attention_windows = attention_windows.view(batch_size, height * width, channels)

        hidden_states = shortcut + self.drop_path(attention_windows)

        layer_output = self.layernorm_after(hidden_states)
        layer_output = self.intermediate(layer_output)
        layer_output = hidden_states + self.output(layer_output)

        outputs = (layer_output,) + outputs

        return outputs


class MaskFormerSwinStage(nn.Module):
    
    def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample):
        super().__init__()
        self.config = config
        self.dim = dim
        self.blocks = nn.ModuleList(
            [
                MaskFormerSwinLayer(
                    config=config,
                    dim=dim,
                    input_resolution=input_resolution,
                    num_heads=num_heads,
                    shift_size=0 if (i % 2 == 0) else config.window_size // 2,
                )
                for i in range(depth)
            ]
        )

        
        if downsample is not None:
            self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm)
        else:
            self.downsample = None

        self.pointing = False

    def forward(
        self, hidden_states, input_dimensions, head_mask=None, output_attentions=False, output_hidden_states=False
    ):
        all_hidden_states = () if output_hidden_states else None

        height, width = input_dimensions
        for i, block_module in enumerate(self.blocks):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None

            block_hidden_states = block_module(hidden_states, input_dimensions, layer_head_mask, output_attentions)

            hidden_states = block_hidden_states[0]

            if output_hidden_states:
                all_hidden_states += (hidden_states,)

        if self.downsample is not None:
            height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2
            output_dimensions = (height, width, height_downsampled, width_downsampled)
            hidden_states = self.downsample(hidden_states, input_dimensions)
        else:
            output_dimensions = (height, width, height, width)

        return hidden_states, output_dimensions, all_hidden_states


class MaskFormerSwinEncoder(nn.Module):
    
    def __init__(self, config, grid_size):
        super().__init__()
        self.num_layers = len(config.depths)
        self.config = config
        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
        self.layers = nn.ModuleList(
            [
                MaskFormerSwinStage(
                    config=config,
                    dim=int(config.embed_dim * 2**i_layer),
                    input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)),
                    depth=config.depths[i_layer],
                    num_heads=config.num_heads[i_layer],
                    drop_path=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
                    downsample=MaskFormerSwinPatchMerging if (i_layer < self.num_layers - 1) else None,
                )
                for i_layer in range(self.num_layers)
            ]
        )

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states,
        input_dimensions,
        head_mask=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_input_dimensions = ()
        all_self_attentions = () if output_attentions else None

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        for i, layer_module in enumerate(self.layers):
            layer_head_mask = head_mask[i] if head_mask is not None else None

            if self.gradient_checkpointing and self.training:
                layer_hidden_states, output_dimensions, layer_all_hidden_states = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    layer_head_mask,
                    output_attentions,
                )
            else:
                layer_hidden_states, output_dimensions, layer_all_hidden_states = layer_module(
                    hidden_states,
                    input_dimensions,
                    layer_head_mask,
                    output_attentions,
                    output_hidden_states,
                )

            input_dimensions = (output_dimensions[-2], output_dimensions[-1])
            all_input_dimensions += (input_dimensions,)
            if output_hidden_states:
                all_hidden_states += (layer_all_hidden_states,)

            hidden_states = layer_hidden_states

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_all_hidden_states[1],)

        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)

        return MaskFormerSwinBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            hidden_states_spatial_dimensions=all_input_dimensions,
            attentions=all_self_attentions,
        )



class MaskFormerSwinPreTrainedModel(PreTrainedModel):
    

    config_class = MaskFormerSwinConfig
    base_model_prefix = "model"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            
            
            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)


class MaskFormerSwinModel(MaskFormerSwinPreTrainedModel):
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config
        self.num_layers = len(config.depths)
        self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))

        self.embeddings = MaskFormerSwinEmbeddings(config)
        self.encoder = MaskFormerSwinEncoder(config, self.embeddings.patch_grid)

        self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
        self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None

    def get_input_embeddings(self):
        return self.embeddings.patch_embeddings

    def _prune_heads(self, heads_to_prune):
        
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        pixel_values=None,
        head_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        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")

        
        
        
        
        
        head_mask = self.get_head_mask(head_mask, len(self.config.depths))

        embedding_output, input_dimensions = self.embeddings(pixel_values)

        encoder_outputs = self.encoder(
            embedding_output,
            input_dimensions,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = encoder_outputs.last_hidden_state if return_dict else encoder_outputs[0]
        sequence_output = self.layernorm(sequence_output)

        pooled_output = None
        if self.pooler is not None:
            pooled_output = self.pooler(sequence_output.transpose(1, 2))
            pooled_output = torch.flatten(pooled_output, 1)

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        hidden_states_spatial_dimensions = (input_dimensions,) + encoder_outputs.hidden_states_spatial_dimensions

        return MaskFormerSwinModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            hidden_states_spatial_dimensions=hidden_states_spatial_dimensions,
            attentions=encoder_outputs.attentions,
        )


class MaskFormerSwinBackbone(MaskFormerSwinPreTrainedModel, BackboneMixin):
    

    def __init__(self, config: MaskFormerSwinConfig):
        super().__init__(config)
        super()._init_backbone(config)

        self.model = MaskFormerSwinModel(config)
        if "stem" in self.out_features:
            raise ValueError("This backbone does not support 'stem' in the `out_features`.")
        self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
        self.hidden_states_norms = nn.ModuleList(
            [nn.LayerNorm(num_channels) for num_channels in self.num_features[1:]]
        )

        
        self.post_init()

    def forward(
        self,
        pixel_values: Tensor,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> BackboneOutput:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        outputs = self.model(
            pixel_values, output_hidden_states=True, output_attentions=output_attentions, return_dict=True
        )

        
        hidden_states = outputs.hidden_states[1:]

        
        
        spatial_dimensions: Tuple[Tuple[int, int]] = outputs.hidden_states_spatial_dimensions
        feature_maps = ()
        for i, (hidden_state, stage, (height, width)) in enumerate(
            zip(hidden_states, self.stage_names[1:], spatial_dimensions)
        ):
            norm = self.hidden_states_norms[i]
            
            hidden_state_unpolled = hidden_state[-1]
            hidden_state_norm = norm(hidden_state_unpolled)
            
            batch_size, _, hidden_size = hidden_state_norm.shape
            
            hidden_state_permuted = (
                hidden_state_norm.permute(0, 2, 1).view((batch_size, hidden_size, height, width)).contiguous()
            )
            if stage in self.out_features:
                feature_maps += (hidden_state_permuted,)

        if not return_dict:
            output = (feature_maps,)
            if output_hidden_states:
                output += (outputs.hidden_states,)
            if output_attentions:
                output += (outputs.attentions,)
            return output

        return BackboneOutput(
            feature_maps=feature_maps,
            hidden_states=outputs.hidden_states if output_hidden_states else None,
            attentions=outputs.attentions,
        )