o Zh @sdZddlZddlZddlZddlmZddlmZm Z m Z m Z ddl Z ddl Z ddl mZmZddlmZmZmZddlmZdd lmZmZmZmZmZmZdd lmZdd lm Z m!Z!m"Z"dd l#m$Z$m%Z%m&Z&dd l'm(Z(ddl)m*Z*e%+e,Z-gdZ.dZ/dZ0eGdddeZ1dSde jde2de3de jfddZ4Gdddej5Z6Gdd d ej5Z7Gd!d"d"ej5Z8Gd#d$d$ej5Z9Gd%d&d&e9Z:Gd'd(d(ej5Z;e9e:d)ZGd.d/d/ej5Z?Gd0d1d1ej5Z@Gd2d3d3ej5ZAGd4d5d5ej5ZBe$Gd6d7d7eZCe$Gd8d9d9eCZDGd:d;d;ej5ZEe$dd?d?eCZFe$d@d=GdAdBdBeCZGGdCdDdDej5ZHGdEdFdFej5ZIGdGdHdHej5ZJGdIdJdJej5ZKGdKdLdLej5ZLe$GdMdNdNeCZMe$dOd=GdPdQdQeCe(ZNgdRZOdS)TzPyTorch BEiT model.N) dataclass)ListOptionalTupleUnion)Tensornn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BackboneOutputBaseModelOutputBaseModelOutputWithPoolingImageClassifierOutputMaskedLMOutputSemanticSegmenterOutput)PreTrainedModel)#compile_compatible_method_lru_cache find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringlogging torch_int) BackboneMixin) BeitConfig)rizmicrosoft/beit-base-patch16-224ztabby, tabby catc@seZdZdZdS)BeitModelOutputWithPoolinga Class for outputs of [`BeitModel`]. 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)`): Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token will be returned. 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. 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. N)__name__ __module__ __qualname____doc__r$r$U/var/www/auris/lib/python3.10/site-packages/transformers/models/beit/modeling_beit.pyr8srFinput drop_probtrainingreturncCsd|dks|s|Sd|}|jdfd|jd}|tj||j|jd}||||}|S)aF 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. r&rrr)dtypedevice)shapendimtorchZrandr,r-Zfloor_div)r'r(r)Z keep_probr.Z random_tensoroutputr$r$r% drop_pathRs r3csTeZdZdZd deeddffdd Zdejdejfdd Z de fd d Z Z S) BeitDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).Nr(r*cst||_dSN)super__init__r()selfr( __class__r$r%r7is  zBeitDropPath.__init__ hidden_statescCst||j|jSr5)r3r(r)r8r;r$r$r%forwardmszBeitDropPath.forwardcCs d|jS)Nzp={})formatr(r8r$r$r% extra_reprps zBeitDropPath.extra_reprr5) r r!r"r#rfloatr7r0rr=strr@ __classcell__r$r$r9r%r4fs r4c sveZdZdZdeddffdd Zdejded edejfd d Z  dd ejd e ej de e dejfddZ ZS)BeitEmbeddingszc Construct the CLS token, position and patch embeddings. Optionally, also the mask token. configr*Ncstttdd|j|_|jr!ttdd|j|_ nd|_ t ||_ |j |_ t |jtjjr8|jn|j|jf|_|j j}|jrUttd|d|j|_nd|_t|j|_dS)Nr)r6r7r Parameterr0zeros hidden_size cls_tokenZuse_mask_token mask_tokenBeitPatchEmbeddingspatch_embeddings patch_size isinstance image_size collectionsabcIterable num_patchesZ use_absolute_position_embeddingsposition_embeddingsDropouthidden_dropout_probdropout)r8rErSr9r$r%r7|s   zBeitEmbeddings.__init__ embeddingsheightwidthc Cs|jdd}|jjdd}tjs||kr||kr|jS|jddddf}|jddddf}|jd}||j} ||j} t|d} |d| | |}|dddd}t j j || | fdd d }|dddd dd|}tj ||fdd S) a This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This method is also adapted to support torch.jit tracing. Adapted from: - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 rNg?rr ZbicubicFsizemode align_cornersdim)r.rTr0Zjit is_tracingrMrreshapepermuter functional interpolateviewcat) r8rXrYrZrSZ num_positionsZclass_pos_embedZpatch_pos_embedrb new_height new_widthZsqrt_num_positionsr$r$r%interpolate_pos_encodings(     z'BeitEmbeddings.interpolate_pos_encoding pixel_valuesbool_masked_posrlcCs|jdur|durtd|j\}}}}||\}\}} |\} } }|durC|j| | d} |d | } |d| | | }|j | dd}t j ||fdd}|jdurb|| |||}||}||| ffS)Nz`interpolate_pos_encoding` argument has no effect for BEiTEmbeddings, embeddings are always interpolated to the input image size. The argument will be removed in transformers v4.51.0.r[rra)rTwarningswarnr.rLr^rJexpand unsqueezeZtype_asrIr0rirlrW)r8rmrnrl_rYrZrX patch_height patch_width batch_sizeZseq_lenZ mask_tokenswZ cls_tokensr$r$r%r=s"   zBeitEmbeddings.forwardNN)r r!r"r#rr7r0rintrlr BoolTensorboolr=rCr$r$r9r%rDvs+rDcs6eZdZdZfddZdejdejfddZZS)rKz 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. cst|j|j}}|j|j}}t|tjj r|n||f}t|tjj r)|n||f}|d|d|d|d}|d|d|d|df}||_||_||_||_ ||_ t j ||||d|_dS)Nrr kernel_sizeZstride)r6r7rOrM num_channelsrHrNrPrQrRrS patch_shaperConv2d projection)r8rErOrMr~rHrSrr9r$r%r7s   zBeitPatchEmbeddings.__init__rmr*c Cs^|j\}}}}||jkrtd||}|jd|jd}}|ddd}|||ffS)NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.r\r r)r.r~ ValueErrorrflatten transpose) r8rmrvr~rYrZrXrtrur$r$r%r=s   zBeitPatchEmbeddings.forward) r r!r"r#r7r0rr=rCr$r$r9r%rKs rKceZdZddedeeddffdd ZddZ   dd ej d eej d e d eej de dee e de e ej e ej ej fffddZZS)BeitSelfAttentionNrE window_sizer*cst||_|j|jdkr"t|ds"td|jd|jd|j|_t|j|j|_|j|j|_ t |j|j |_ t j |j|j dd|_ t |j|j |_t |j|_t||_|jrkt||d|_dSdS) NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .F)biasr)r6r7rErHnum_attention_headshasattrrryattention_head_size all_head_sizerLinearquerykeyvaluerUattention_probs_dropout_probrWr{has_relative_position_biasBeitRelativePositionBiasrelative_position_biasr8rErr9r$r%r7s&   zBeitSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)Nr[rr\rr )r^rrrhre)r8xZ new_x_shaper$r$r%transpose_for_scoress z&BeitSelfAttention.transpose_for_scoresFr; head_maskoutput_attentionsrrl resolutioncCs.||}|||}|||} ||} t| |dd} | t|j } |j rL|\} } | |j j | |j j f}| |j |||jdd} |durT| |} tjj| dd}||}|duri||}t|| }|dddd}|dd|jf}|j|}|r||f}|S|f}|S) Nr[rdim_sizerarr\r )rrrrr0matmulrmathsqrtrrrErMrr.rrfZsoftmaxrWre contiguousr^rrh)r8r;rrrrlrmixed_query_layer key_layer value_layer query_layerZattention_scoresrYrZrZattention_probs context_layernew_context_layer_shapeoutputsr$r$r%r=s4       zBeitSelfAttention.forwardr5NFNFN)r r!r"rrtupler7rr0rr{rryrr=rCr$r$r9r%rs.  rcsveZdZ     d dejdeejdedeejdedeeed e eejeejejffffd d Z Z S) BeitSdpaSelfAttentionNFr;rrrrlrr*c s8|s|durtdtj||||||dS||}|||}|||} ||} d} |jrQ|\} } | |j j | |j j f}|j |||j dd} |dur`| dur\|} n| |7} dt |j}tjjj| || | |jrw|j jndd|d}|dd dd }|dd |jf}|j|}|dfS) Na`BeitSdpaSelfAttention` is used but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `head_mask`. Falling back to the manual attention implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.)r;rrrrlrrrr&F)Z attn_maskZ dropout_pZ is_causalscalerr\r r)loggerZ warning_oncer6r=rrrrrrErMrr.rrrr0rrfZscaled_dot_product_attentionr)rrerr^rrh)r8r;rrrrlrrrrrZ attn_biasrYrZrZscalingrrr9r$r%r=VsR      zBeitSdpaSelfAttention.forwardr) r r!r"r0rrr{rryrr=rCr$r$r9r%rUs* rcsHeZdZdZdeddffdd Zd dejdejdejfd d ZZ S) BeitSelfOutputz The residual connection is defined in BeitLayer instead of here (as is the case with other models), due to the layernorm applied before each block. rEr*Ncs.tt|j|j|_t|j|_dSr5) r6r7rrrHdenserUrVrWr8rEr9r$r%r7 zBeitSelfOutput.__init__r; input_tensorcC||}||}|Sr5rrW)r8r;rgammar$r$r%r=  zBeitSelfOutput.forwardr5) r r!r"r#rr7r0rr=rCr$r$r9r%rs&r)eagerZsdpacr) BeitAttentionNrErr*cs4tt|j||d|_t||_t|_dS)Nr) r6r7BEIT_SELF_ATTENTION_CLASSESZ_attn_implementation attentionrr2set pruned_headsrr9r$r%r7s   zBeitAttention.__init__cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrra)lenrrrrrrrrrr2rrunion)r8headsindexr$r$r% prune_headss zBeitAttention.prune_headsFr;rrrrlrc Cs:|||||||}||d|}|f|dd} | S)Nrr)rr2) r8r;rrrrlrZ self_outputsattention_outputrr$r$r%r=s  zBeitAttention.forwardr5r)r r!r"rrrr7rr0rr{rryrr=rCr$r$r9r%rs.  rc<eZdZdeddffdd ZdejdejfddZZS) BeitIntermediaterEr*NcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr5) r6r7rrrHintermediate_sizerrNZ hidden_actrBr intermediate_act_fnrr9r$r%r7s   zBeitIntermediate.__init__r;cCrr5)rrr<r$r$r%r=rzBeitIntermediate.forward r r!r"rr7r0rr=rCr$r$r9r%rsrcr) BeitOutputrEr*Ncs.tt|j|j|_t|j|_ dSr5) r6r7rrrrHrrUrVrWrr9r$r%r7rzBeitOutput.__init__r;cCrr5rr<r$r$r%r=rzBeitOutput.forwardrr$r$r9r%rsrcseZdZdZddedeededdffdd Z   dd e j d ee j d e dee j de dee e dee e j e e j e j fffddZZS) BeitLayerz?This corresponds to the Block class in the timm implementation.Nr&rErdrop_path_rater*cst|j|_d|_t||d|_t||_t||_ t j |j |j d|_|dkr/t|nt |_t j |j |j d|_|j}|dkrct j|t|j dd|_t j|t|j dd|_dSd\|_|_dS) NrrZepsr&rT)Z requires_gradrx)r6r7Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterr2r LayerNormrHlayer_norm_epslayernorm_beforer4Identityr3layernorm_afterlayer_scale_init_valuerFr0Zoneslambda_1lambda_2)r8rErrZ init_valuesr9r$r%r7s    zBeitLayer.__init__Fr;rrrrlrc Cs|j|||||||d}|d}|dd} |jdur"|j|}|||}||} || } || } |jdurB|j| } || |} | f| } | S)N)rrrlrrr)rrrr3rrr2r) r8r;rrrrlrZself_attention_outputsrrZ layer_outputr$r$r%r= s*          zBeitLayer.forward)Nr&r)r r!r"r#rrrrAr7r0rr{rryrr=rCr$r$r9r%rs.$ rcsfeZdZdededdffdd Zedddeeefde j fd d Z dd e de j fd dZ ZS)rrErr*NcsRt||_d|ddd|ddd|_tt|j|j|_ dS)Nr\rrr ) r6r7rnum_relative_distancerrFr0rGrrelative_position_bias_tablerr9r$r%r76s & z!BeitRelativePositionBias.__init__ )maxsizec Csd|ddd|ddd}|d|d}tjt|dt|ddd}t|}t|d}|dddddf|dddddf}|ddd}|dddddf|dd7<|dddddf|dd7<|dddddfd|dd9<tj|dfd|jd}| d |ddddf<|d|dddf<|d|dddf<|d|d <|S) z This method creates the relative position index, modified to support arbitrary window sizes, as introduced in [MiDaS v3.1](https://arxiv.org/abs/2307.14460). r\rrr Zij)ZindexingN)r^r,r[)rr) r0ZmeshgridZarangestackrrerrGr,sum) r8rrZ window_areagridZcoordsZcoords_flattenZrelative_coordsrelative_position_indexr$r$r% generate_relative_position_index?s $$  ,&&* z9BeitRelativePositionBias.generate_relative_position_indexFrlcCsfd|jdd}d|jdd}d|dd}d|dd}|j}|j} ||d} |d| d} | d||ddddd} tjj| t|t|fdd} | dddd| dd} t | || ddg} | |}| | d}| |d|dd|d|ddd}|ddd }|rtjj|d||fdd d d}|dS) zu Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes. r\rrr Nr[bilinear)r^r_Fr])rrrrdrerrfrgrr0rirrhrrrsqueeze)r8rrlrZ old_heightZ old_widthrjrkZ old_relative_position_bias_tableZold_num_relative_distanceZnew_num_relative_distanceZ old_sub_tableZ new_sub_tableZ new_relative_position_bias_tablerrr$r$r%r=Xs@  & z BeitRelativePositionBias.forward)FN)r r!r"rrr7rrryr0rrr{r=rCr$r$r9r%r5s  rcseZdZddedeeddffdd Z      dd ejd eejd e d e d e dee e e fde de ee ffddZZS) BeitEncoderNrErr*cszt|_j|_|jrtd|_ddtjdj j ddDt fddt j D|_d|_dS) NrcSsg|]}|qSr$)item.0rr$r$r% sz(BeitEncoder.__init__..rcpu)r-cs(g|]}tjr nd|dqS)N)rr)rZuse_relative_position_biasrirEZdprrr$r%rs F)r6r7rEZ!use_shared_relative_position_biasrrrr0Zlinspacernum_hidden_layersr ModuleListrangelayergradient_checkpointingrr9rr%r7s   zBeitEncoder.__init__FTr;rroutput_hidden_statesrlr return_dictc Cs|rdnd}|r dnd} t|jD]_\} } |r||f}|jr;|\} } | |jj| |jjf}|j|||jdd}nd}|durE|| nd}|jrZ|jrZ| | j ||||||}n | ||||||}|d}|rp| |df} q|rx||f}|st dd||| fDSt ||| dS)Nr$r)rlrrcss|] }|dur|VqdSr5r$)rvr$r$r% sz&BeitEncoder.forward..)last_hidden_stater; attentions) enumeraterrrErMrr.rr)Z_gradient_checkpointing_func__call__rr)r8r;rrrrlrrZall_hidden_statesZall_self_attentionsrZ layer_modulerYrZrrZlayer_head_maskZ layer_outputsr$r$r%r=sX       zBeitEncoder.forwardr5)NFFFNT)r r!r"rrrr7r0rr{rryrrr=rCr$r$r9r%rs2  rc@s4eZdZeZdZdZdZdgZdgZ dZ ddZ dS) BeitPreTrainedModelbeitrmTrz.*relative_position_index.*cCsbt|tjtjtjfr%|jjjd|jj d|j dur#|j j dSdSt|tj rH|jjjd|jj d|j durF|jj|j  dSdSt|tjr]|j j |jjddSt|tr|jj |jdurs|jj |jdur|jj dSdSt|tr|jj dSt|tr|jdur|jj|jj|jj|jjdSdSdS)zInitialize the weightsr&)meanZstdNg?)rNrrrConvTranspose2dweightdataZnormal_rEZinitializer_rangerZzero_Z EmbeddingZ padding_idxrZfill_rDrIrJrTrrrrrr)r8moduler$r$r% _init_weightss8             z!BeitPreTrainedModel._init_weightsN) r r!r"rZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesZ"_keys_to_ignore_on_load_unexpectedZ_supports_sdparr$r$r$r%rs rcseZdZddededdffdd Zdd Zd d Ze     dd e j de e j de e j de ede edede ede eeffddZZS) BeitModelTrEadd_pooling_layerr*Ncspt|||_t||_t||jjjd|_|j rt nt j |j |jd|_|r/t|nd|_|dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer rrN)r6r7rErDrXrrLrencoderuse_mean_poolingrrrrHr layernorm BeitPoolerpooler post_init)r8rErr9r$r%r7s   zBeitModel.__init__cC|jjSr5rXrLr?r$r$r%get_input_embeddingszBeitModel.get_input_embeddingscCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsrrrr)r8Zheads_to_prunerrr$r$r% _prune_headsszBeitModel._prune_headsFrmrnrrrrlrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|||jj}|j||d\}} |jdd} |j||||| ||d} | d} | | } |j durU| | nd} |sl| dura| | fn| f}|| ddSt | | | j | j dS)z bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). N)rnr\)rrrrrrlrr)r pooler_outputr;r)rErruse_return_dictZ get_head_maskrrXr.rrrrr;r)r8rmrnrrrrlrembedding_outputrsrZencoder_outputssequence_output pooled_outputZ head_outputsr$r$r%r=&s8 zBeitModel.forward)T)NNNNFN)r r!r"rr{r7r r rr0rrrzrrrr=rCr$r$r9r%rs8 rcr) rrEr*Ncs2t|jrtj|j|jd|_dSd|_dS)Nr)r6r7rrrrHrrrr9r$r%r7_s zBeitPooler.__init__r;cCsL|jdur|ddddddf}||d}|S|dddf}|S)Nrr)rr)r8r;Z patch_tokensrr$r$r%r=es zBeitPooler.forwardrr$r$r9r%r^sra Beit Model transformer with a 'language' modeling head on top. BEiT does masked image modeling by predicting visual tokens of a Vector-Quantize Variational Autoencoder (VQ-VAE), whereas other vision models like ViT and DeiT predict RGB pixel values. As a result, this class is incompatible with [`AutoModelForMaskedImageModeling`], so you will need to use [`BeitForMaskedImageModeling`] directly if you wish to do masked image modeling with BEiT. )Z custom_introcseZdZdeddffdd Ze        ddeejdeej d eejd eejd ee d ee d e dee de e e ffddZZS)BeitForMaskedImageModelingrEr*NcsTt||j|_t|dd|_tj|j|jd|_ t |j|j |_ | dS)NFrr)r6r7 num_labelsrrrrrHrrrZ vocab_sizelm_headrrr9r$r%r7zs  z#BeitForMaskedImageModeling.__init__Frmrnrlabelsrrrlrc  Cs|dur|n|jj}|j|||||||d} | d} || } || ddddf} d} |dur 1` a classification loss is computed (Cross-Entropy). Examples: ```python >>> from transformers import AutoImageProcessor, BeitForMaskedImageModeling >>> import torch >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224-pt22k") >>> model = BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k") >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2 >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values >>> # create random boolean mask of shape (batch_size, num_patches) >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool() >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) >>> loss, logits = outputs.loss, outputs.logits >>> list(logits.shape) [1, 196, 8192] ```N)rnrrrrlrrrlosslogitsr;r) rErrrrr rr;r)r8rmrnrrrrrlrrrZprediction_scoresZmasked_lm_lossloss_fctr2r$r$r%r=s4, z"BeitForMaskedImageModeling.forward)NNNNNNFN)r r!r"rr7rrr0rrzr{rrrr=rCr$r$r9r%rqs<    rz Beit Model transformer with an image classification head on top (a linear layer on top of the average of the final hidden states of the patch tokens) e.g. for ImageNet. cseZdZdeddffdd Ze       ddeejdeejd eejd ee d ee d e d ee de e e ffddZ ZS)BeitForImageClassificationrEr*NcsRt||j|_t|dd|_|jdkrt|j|jnt|_ | dS)NTrr) r6r7rrrrrrHr classifierrrr9r$r%r7s $ z#BeitForImageClassification.__init__FrmrrrrrlrcCsl|dur|n|jj}|j||||||d}|r|jn|d} || } d} |dur|jjdurS|jdkr9d|j_n|jdkrO|jtj ksJ|jtj krOd|j_nd|j_|jjdkrqt } |jdkrk| | | } n+| | |} n%|jjdkrt } | | d|j|d} n|jjdkrt} | | |} |s| f|dd} | dur| f| S| St| | |j|jd S) a 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). NrrrrlrrZ regressionZsingle_label_classificationZmulti_label_classificationr[r\r)rErrrrZ problem_typerr,r0longryr rr rhr rr;r)r8rmrrrrrlrrrrrrr2r$r$r%r=sN    "       z"BeitForImageClassification.forwardNNNNNFN)r r!r"rr7rrr0rr{rrrr=rCr$r$r9r%rs6  rcseZdZdZ   ddededeeeeeffdeeeeefefd ed eeeeeffd d ffd d Z de j d e j fddZ Z S)BeitConvModuleaD A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU). Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. rFr in_channels out_channelsr}paddingrdilationr*Ncs<ttj||||||d|_t||_t|_dS)N)r"r#r}r$rr%) r6r7rrconv BatchNorm2dbnZReLU activation)r8r"r#r}r$rr%r9r$r%r70s  zBeitConvModule.__init__r'cCs"||}||}||}|Sr5)r&r(r))r8r'r2r$r$r%r=Es   zBeitConvModule.forward)rFr)r r!r"r#ryrrrBr{r7r0rr=rCr$r$r9r%r!(s* r!csDeZdZdedededdffdd Zdejdejfd d ZZS) BeitPyramidPoolingBlock pool_scaler"channelsr*NcsLtt|t||ddg|_t|jD] \}}|t||qdS)Nrr}) r6r7rZAdaptiveAvgPool2dr!layersr add_modulerB)r8r+r"r,rrr9r$r%r7Ns  z BeitPyramidPoolingBlock.__init__r'cCs|}|jD]}||}q|Sr5)r.)r8r' hidden_staterr$r$r%r=Ws  zBeitPyramidPoolingBlock.forward) r r!r"ryr7r0rr=rCr$r$r9r%r*Ms r*c sXeZdZdZdeedfdedededdf fd d Zd ej de ej fd d Z Z S)BeitPyramidPoolingModulea Pyramid Pooling Module (PPM) used in PSPNet. Args: pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid Module. in_channels (int): Input channels. channels (int): Channels after modules, before conv_seg. align_corners (bool): align_corners argument of F.interpolate. Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. pool_scales.r"r,r`r*Ncsht||_||_||_||_g|_t|D]\}}t|||d}|j || t ||qdS)N)r+r"r,) r6r7r2r`r"r,blocksrr*appendr/rB)r8r2r"r,r`rr+blockr9r$r%r7ls  z!BeitPyramidPoolingModule.__init__rcCsHg}|jD]}||}tjj||ddd|jd}||q|S)Nr\rr])r3rrfrgr^r`r4)r8rZppm_outsppmZppm_outZupsampled_ppm_outr$r$r%r=xs  z BeitPyramidPoolingModule.forward) r r!r"r#rryr{r7r0rrr=rCr$r$r9r%r1^s* " r1csHeZdZdZdeddffdd ZddZd ejdejfd d Z Z S) BeitUperHeadz Unified Perceptual Parsing for Scene Understanding. This head is the implementation of [UPerNet](https://arxiv.org/abs/1807.10221). Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. rEr*Ncst|j|_|jgd|_|j|_d|_tj|j|j dd|_ t |j|jd|j|jd|_ t |jdt|j|j|jddd|_t|_t|_|jddD] }t ||jdd}t |j|jddd}|j||j|qWt t|j|j|jddd|_dS) NFrr-r[)r`r r}r$)r6r7r2rHr"r,r`rrrrr1 psp_modulesr!r bottleneckr lateral_convs fpn_convsr4fpn_bottleneck)r8rEr"Zl_convZfpn_convr9r$r%r7s>     zBeitUperHead.__init__cCs:|d}|g}|||tj|dd}||}|S)Nr[rra)extendr:r0rir;)r8ZinputsrZpsp_outsr2r$r$r% psp_forwards  zBeitUperHead.psp_forwardencoder_hidden_statescsfddtjDt}t|dddD]$}|djdd}|dtjj ||dj d|d<q fd dt|dD}|dt|dddD]}tjj |||djdddj d||<qbt j |dd } |}|}|S) Ncsg|] \}}||qSr$r$)rrZ lateral_conv)rAr$r%rsz(BeitUperHead.forward..rrr[r\rr]csg|] }j||qSr$)r=r)lateralsr8r$r%rsra)rr<r4r@rrr.rrfrgr`r0rir>r)r8rAZused_backbone_levelsrZ prev_shapeZfpn_outsr2r$)rArBr8r%r=s$   zBeitUperHead.forward) r r!r"r#rr7r@r0rr=rCr$r$r9r%r7s & r7c s`eZdZdZ ddedededeeeeeffd d f fd d Zd e j d e j fddZ Z S) BeitFCNHeada Fully Convolution Networks for Semantic Segmentation. This head is implemented of [FCNNet](https://arxiv.org/abs/1411.4038>). Args: config (BeitConfig): Configuration. in_channels kernel_size (int): The kernel size for convs in the head. Default: 3. dilation (int): The dilation rate for convs in the head. Default: 1. Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. r\r rrEin_indexr}r%r*Nc st|j|_|j|_|j|_|j|_ ||_ |d|}g}| t |j|j|||dt |jdD]}| t |j|j|||dq5|jdkrQt|_ntj||_|j rjt |j|j|j||dd|_tj|j|jdd|_dS)Nr\)r}r$r%rrr9r-)r6r7rHr"Zauxiliary_channelsr,Zauxiliary_num_convsZ num_convsZauxiliary_concat_input concat_inputrDr4r!rrrconvs Sequentialconv_catrrr)r8rErDr}r%Z conv_paddingrFrr9r$r%r7s6     zBeitFCNHead.__init__rAcCs@||j}||}|jr|tj||gdd}||}|S)Nrra)rDrFrErHr0rir)r8rAr;r2r$r$r%r= s   zBeitFCNHead.forward)r\r r) r r!r"r#rryrrr7r0rr=rCr$r$r9r%rCs"rCcseZdZdeddffdd ZddZe       dd eej d eej d eej d ee d ee de dee de e e ffddZZS)BeitForSemanticSegmentationrEr*Nc st||j|_t|dd|_t|jjdkrtdt t j |j |j dddt |j t t j |j |j ddd|_t t j |j |j ddd|_t |_t jddd|_t||_|jrft|nd|_|dS)NFrr8zBeitForSemanticSegmentation requires config.out_indices to be a list of 4 integers, specifying which features to use from the backbone. One can use [3, 5, 7, 11] in case of a base-sized architecture.r\r|)r6r7rrrrrE out_indicesrrrGrrHr'GELUfpn1fpn2rfpn3 MaxPool2dfpn4r7 decode_headZuse_auxiliary_headrCauxiliary_headrrr9r$r%r7s*     z$BeitForSemanticSegmentation.__init__c Cstjj||jddddd}|dur"tjj||jddddd}t|jjd}|||}|}|durA|||} ||jj| 7}|S)NrrFr])Z ignore_index)rrfrgr.r rEZsemantic_loss_ignore_indexZauxiliary_loss_weight) r8rauxiliary_logitsrZupsampled_logitsZupsampled_auxiliary_logitsrZ main_lossrZauxiliary_lossr$r$r% compute_loss5s  z(BeitForSemanticSegmentation.compute_lossFrmrrrrrlrcs|dur|njj}|dur|njj}|dur"jjdkr"tdj|||d||d}|r2|jn|d} fddt| D} |jdjj jj fd d| D} j j j jg} tt| D] } | | | | | | <qh| } d}jdur| }d}|dur| ||}|s|r| f|dd}n | f|d d}|dur|f|S|St|| |r|jnd|jd S) aD labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*): Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy). Examples: ```python >>> from transformers import AutoImageProcessor, BeitForSemanticSegmentation >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-finetuned-ade-640-640") >>> model = BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> # logits are of shape (batch_size, num_labels, height, width) >>> logits = outputs.logits ```Nrz/The number of labels should be greater than oneTrcs$g|]\}}|djjvr|qSr+)rErJ)ridxfeaturer?r$r%rs$z7BeitForSemanticSegmentation.forward..rcs<g|]}|ddddddfddddqS)Nrrr\r[)rerdr)rvpatch_resolutionr$r%rs0r\r)rErrrrrr;rr.rOrMrLrMrNrPrrrQrRrTrr)r8rmrrrrrlrrrAfeaturesopsrrrSrr2r$)rvrWr8r%r=HsR"       z#BeitForSemanticSegmentation.forwardr )r r!r"rr7rTrrr0rr{rrrr=rCr$r$r9r%rIs8  rIzM BEiT backbone, to be used with frameworks like DETR and MaskFormer. c sZeZdZfddZddZe   d dedeedeed eed e f d d Z Z S) BeitBackbonec sttfddtjdD|_t|_t|jj j d|_ j rut |jjdkr9tdj}ttj||dddtj|jd ttj||ddd|_ttj||ddd|_t|_tjddd|_|dS) Ncsg|]}jqSr$)rH)rrsrEr$r%rsz)BeitBackbone.__init__..rrr8zBeitBackbone requires config.out_indices to be a list of 4 integers, specifying which features to use from the backbone. One can use [3, 5, 7, 11] in case of a base-sized architecture.r\r|r)r6r7Z_init_backbonerrZ num_featuresrDrXrrLrradd_fpnrrErJrrHrrGrr'Zbatch_norm_epsrKrLrMrrNrOrPr)r8rErHr9r[r%r7s*     zBeitBackbone.__init__cCrr5r r?r$r$r%r r z!BeitBackbone.get_input_embeddingsNrmrrrr*cCs|dur|n|jj}|dur|n|jj}|dur|n|jj}|jd}||\}\}}|jdd} |j|d|| |d} |rB| jn| d} d} t|j | D].\} }| |j vr||jj rw|ddddddf}| ddd}| |d||}| |f7} qN|jjr|| d|| d|| d|| d g} t| } |s|r| f| dd}|S| f| dd}|St| |r| jnd| jd S) a: Examples: ```python >>> from transformers import AutoImageProcessor, AutoBackbone >>> import torch >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224") >>> model = AutoBackbone.from_pretrained( ... "microsoft/beit-base-patch16-224", out_features=["stage1", "stage2", "stage3", "stage4"] ... ) >>> inputs = processor(image, return_tensors="pt") >>> outputs = model(**inputs) >>> feature_maps = outputs.feature_maps >>> list(feature_maps[-1].shape) [1, 768, 14, 14] ```Nrr\T)rrrrrr$r[r ) feature_mapsr;r)rErrrr.rXrr;zipZ stage_namesZ out_featuresZreshape_hidden_statesrerdr\rLrMrNrPrrr)r8rmrrrrvrrtrurrr;r]Zstager0r2r$r$r%r=sT         zBeitBackbone.forward)NNN) r r!r"r7r rrrr{rr=rCr$r$r9r%rZs$ rZ)rrrIrrrZ)r&F)Pr#collections.abcrPrro dataclassesrtypingrrrrr0Ztorch.utils.checkpointrrZtorch.nnr r r Z activationsr Zmodeling_outputsrrrrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrZutils.backbone_utilsrZconfiguration_beitrZ get_loggerr rZ_EXPECTED_OUTPUT_SHAPEZ_IMAGE_CLASS_CHECKPOINTZ_IMAGE_CLASS_EXPECTED_OUTPUTrrAr{r3Moduler4rDrKrrrrrrrrrrrrrrrr!r*r1r7rCrIrZ__all__r$r$r$r%s        f&S>, ASW&W[N%%U; w