a h@sdZddlZddlZddlZddlmZddlmZm Z ddl Z ddl Z ddl m Z m Z ddlmZmZmZddlmZdd lmZdd lmZmZmZmZmZmZdd lmZdd lm Z m!Z!m"Z"dd l#m$Z$m%Z%m&Z&ddl'm(Z(ddl)m*Z*e%+e,Z-ee$ddGdddeZ.dOe j e/e0e j dddZ1Gddde j2Z3Gddde j2Z4Gddde j2Z5Gd d!d!e j2Z6Gd"d#d#e6Z7Gd$d%d%e j2Z8e6e7d&Z9Gd'd(d(e j2Z:Gd)d*d*e j2Z;Gd+d,d,e j2ZGd1d2d2e j2Z?e$Gd3d4d4eZ@e$Gd5d6d6e@ZAGd7d8d8e j2ZBe$d9dGd:d;d;e@ZCe$dd>e@ZDGd?d@d@e j2ZEGdAdBdBe j2ZFGdCdDdDe j2ZGGdEdFdFe j2ZHGdGdHdHe j2ZIe$GdIdJdJe@ZJe$dKdGdLdMdMe@e(ZKgdNZLdS)PzPyTorch BEiT model.N) dataclass)OptionalUnion)Tensornn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)GradientCheckpointingLayer)BackboneOutputBaseModelOutputBaseModelOutputWithPoolingImageClassifierOutputMaskedLMOutputSemanticSegmenterOutput)PreTrainedModel)#compile_compatible_method_lru_cache find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringlogging torch_int) BackboneMixin) BeitConfigz- Class for outputs of [`BeitModel`]. )Z custom_introc@seZdZdZdS)BeitModelOutputWithPoolingaF 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. N)__name__ __module__ __qualname____doc__r"r"b/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/beit/modeling_beit.pyr0srF)input 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_path?s r1csPeZdZdZd eeddfdd ZejejdddZ e d d d Z Z S) BeitDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).N)r&r(cst||_dSN)super__init__r&)selfr& __class__r"r#r5Vs zBeitDropPath.__init__ hidden_statesr(cCst||j|jSr3)r1r&r'r6r:r"r"r#forwardZszBeitDropPath.forward)r(cCs d|jS)Nzp=)r&r6r"r"r# extra_repr]szBeitDropPath.extra_repr)N) rrr r!rfloatr5r.rr<strr> __classcell__r"r"r7r#r2Ssr2csdeZdZdZeddfdd ZejeeejdddZ d eje ej e e ejd d d Z ZS) BeitEmbeddingszc Construct the CLS token, position and patch embeddings. Optionally, also the mask token. Nconfigr(cstttdd|j|_|jrBttdd|j|_ nd|_ t ||_ |j |_ t |jtjjrp|jn |j|jf|_|j j}|jrttd|d|j|_nd|_t|j|_dS)Nr)r4r5r Parameterr.zeros hidden_size cls_tokenZuse_mask_token mask_tokenBeitPatchEmbeddingspatch_embeddings patch_size isinstance image_size collectionsabcIterable num_patchesZ use_absolute_position_embeddingsposition_embeddingsDropouthidden_dropout_probdropout)r6rDrRr7r"r#r5is    zBeitEmbeddings.__init__) embeddingsheightwidthr(c 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,rSr.Zjit is_tracingrLrreshapepermuter functional interpolateviewcat) r6rWrXrYrRZ num_positionsZclass_pos_embedZpatch_pos_embedra new_height new_widthZsqrt_num_positionsr"r"r#interpolate_pos_encodings(     z'BeitEmbeddings.interpolate_pos_encoding) pixel_valuesbool_masked_posrkr(cCs|jdur|durtd|j\}}}}||\}\}} |\} } }|dur|j| | d} |d | } |d| | | }|j | dd}t j ||fdd}|jdur|| |||}||}||| 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.rZrr`)rSwarningswarnr,rKr]rIexpand unsqueezeZtype_asrHr.rhrkrV)r6rlrmrk_rXrYrW patch_height patch_width batch_sizeZseq_lenZ mask_tokenswZ cls_tokensr"r"r#r<s"  zBeitEmbeddings.forward)NN)rrr r!rr5r.rintrkr BoolTensorboolr<rAr"r"r7r#rBcs+rBcs4eZdZdZfddZejejdddZZS)rJz 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 r8|n||f}t|tjj rR|n||f}|d|d|d|d}|d|d|d|df}||_||_||_||_ ||_ t j ||||d|_dS)Nrr kernel_sizeZstride)r4r5rNrL num_channelsrGrMrOrPrQrR patch_shaperConv2d projection)r6rDrNrLr|rGrRr}r7r"r#r5s   zBeitPatchEmbeddings.__init__)rlr(c Cs^|j\}}}}||jkr td||}|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) r6rlrur|rXrYrWrsrtr"r"r#r<s  zBeitPatchEmbeddings.forward) rrr r!r5r.rr<rAr"r"r7r#rJs rJc sxeZdZd eeeddfdd Zd ejeeje eeje eee e eejeejejffdddZ Z S) BeitSelfAttentionNrD window_sizer(cst||_|j|jdkrDt|dsDtd|jd|jd|j|_t|j|j|_|j|j|_ t |j|j |_ t j |j|j dd|_ t |j|j |_t |j|_t||_|jrt||d|_dS) NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .F)biasr)r4r5rDrGnum_attention_headshasattrrrwattention_head_size all_head_sizerLinearquerykeyvaluerTattention_probs_dropout_probrVryhas_relative_position_biasBeitRelativePositionBiasrelative_position_biasr6rDrr7r"r#r5s$   zBeitSelfAttention.__init__Fr: head_maskoutput_attentionsrrk resolutionr(cCsl|j\}}} |||d|j|jdd} |||d|j|jdd} |||d|j|jdd} t | | dd} | t |j} |j r|\}}||j j||j jf}| |j|||jdd} |dur| |} tjj| dd}||}|dur||}t || }|dddd}|dd|jf}|j|}|rb||fn|f}|S) NrZrr[dim_sizer`rr )r,rrgrrrrrr.matmulmathsqrtrrDrLrrreZsoftmaxrVrd contiguousr]r)r6r:rrrrkrru seq_lengthrr query_layer key_layer value_layerZattention_scoresrXrYrZattention_probs context_layernew_context_layer_shapeoutputsr"r"r#r<sN            zBeitSelfAttention.forward)N)NFNFN)rrr rrtupler5r.rryrwrr<rAr"r"r7r#rs rc s`eZdZdejeejeeejeeeee eejeejejffdfdd Z Z S)BeitSdpaSelfAttentionNFrc sx|s |dur.tdtj||||||dS|j\}}} |||d|j|j dd} | ||d|j|j dd} | ||d|j|j dd} d} |j r|\}}||j j||j jf}|j|||jdd} |dur| dur|} n| |7} dt|j}tjjj| | | | |jr.|j jndd|d }|d ddd }|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:rrrrkrrZrr[rr$F)Z attn_maskZ dropout_pZ is_causalscalerr r)loggerZ warning_oncer4r<r,rrgrrrrrrrDrLrrrr.rreZscaled_dot_product_attentionr'rrdrr]r)r6r:rrrrkrrurrrrrrZ attn_biasrXrYrZscalingrrr7r"r#r<Isp           zBeitSdpaSelfAttention.forward)NFNFN) rrr r.rrryrrwrr<rAr"r"r7r#rHs rcsBeZdZdZeddfdd Zd ejejejdddZZ 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. NrCcs.tt|j|j|_t|j|_dSr3) r4r5rrrGdenserTrUrVr6rDr7r"r#r5s zBeitSelfOutput.__init__)r: input_tensorr(cCs||}||}|Sr3rrV)r6r:rgammar"r"r#r<s  zBeitSelfOutput.forward)N) rrr r!rr5r.rr<rAr"r"r7r#rsr)eagerZsdpac seZdZd eeeddfdd ZddZd ej eej e eej e eee e eej eej ej ffdd d Z ZS) BeitAttentionNrcs4tt|j||d|_t||_t|_dS)Nr) r4r5BEIT_SELF_ATTENTION_CLASSESZ_attn_implementation attentionrr0set pruned_headsrr7r"r#r5s  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)Nrrr`)lenrrrrrrrrrr0rrunion)r6headsindexr"r"r# prune_headss zBeitAttention.prune_headsFrc Cs:|||||||}||d|}|f|dd} | S)Nrr)rr0) r6r:rrrrkrZ self_outputsattention_outputrr"r"r#r<s  zBeitAttention.forward)N)NFNFN)rrr rrrr5rr.rryrwrr<rAr"r"r7r#rs  rcs8eZdZeddfdd ZejejdddZZS)BeitIntermediateNrCcsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr3) r4r5rrrGintermediate_sizerrMZ hidden_actr@r intermediate_act_fnrr7r"r#r5s   zBeitIntermediate.__init__r9cCs||}||}|Sr3)rrr;r"r"r#r<s  zBeitIntermediate.forward rrr rr5r.rr<rAr"r"r7r#rsrcs8eZdZeddfdd ZejejdddZZS) BeitOutputNrCcs.tt|j|j|_t|j|_ dSr3) r4r5rrrrGrrTrUrVrr7r"r#r5s zBeitOutput.__init__r9cCs||}||}|Sr3rr;r"r"r#r<s  zBeitOutput.forwardrr"r"r7r#rsrc seZdZdZd eeeeddfdd Zd e j ee j e ee j e eee e fe ee j ee j e j ffdd d ZZS) BeitLayerz?This corresponds to the Block class in the timm implementation.Nr$)rDrdrop_path_rater(cst|j|_d|_t||d|_t||_t||_ t j |j |j d|_|dkr^t|nt |_t j |j |j d|_|j}|dkrt j|t|j dd|_t j|t|j dd|_n d\|_|_dS) Nrrepsr$rT)Z requires_grad)NN)r4r5Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterr0r LayerNormrGlayer_norm_epslayernorm_beforer2Identityr1layernorm_afterlayer_scale_init_valuerEr.Zoneslambda_1lambda_2)r6rDrrZ init_valuesr7r"r#r5s   zBeitLayer.__init__Frc Cs|j|||||||d}|d}|dd} |jdurD|j|}|||}||} || } || } |jdur|j| } || |} | f| } | S)N)rrrkrrr)rrrr1rrr0r) r6r:rrrrkrZself_attention_outputsrrZ layer_outputr"r"r#r<s*          zBeitLayer.forward)Nr$)NFNFN)rrr r!rrrr?r5r.rryrwrr<rAr"r"r7r#rs rcs^eZdZeeddfdd Zeddeeefej ddd Z de ej d d d Z Z S)rNrcsRt||_d|ddd|ddd|_tt|j|j|_ dS)Nr[rrr ) r4r5rnum_relative_distancerrEr.rFrrelative_position_bias_tablerr7r"r#r55s  &z!BeitRelativePositionBias.__init__ )maxsize)rr(c 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://huggingface.co/papers/2307.14460). r[rrr Zij)ZindexingN)r]r*rZ)rr) r.ZmeshgridZarangestackrrdrrFr*sum) r6rrZ window_areagridZcoordsZcoords_flattenZrelative_coordsrelative_position_indexr"r"r# generate_relative_position_index>s $$  ,&&* z9BeitRelativePositionBias.generate_relative_position_indexF)rkr(cCshd|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 }|r^tjj|d||fdd d d}|dS) zu Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes. r[rrr NrZbilinear)r]r^Fr\)rrrrcrdrrerfrr.rhrrgrrqsqueeze)r6rrkrZ old_heightZ old_widthrirjZ 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<Ws@  &z BeitRelativePositionBias.forward)FN)rrr rrr5rrwr.rrryr<rAr"r"r7r#r4s rc sfeZdZd eeeddfdd Zd ejeeje e e eee e fe e ee fddd Z ZS) BeitEncoderNrcszt|_j|_|jr,td|_ddtjdj j ddDt fddt j D|_d|_dS) NrcSsg|] }|qSr")item.0xr"r"r# z(BeitEncoder.__init__..rcpu)r+cs(g|] }tjrnd|dqS)N)rr)rZuse_relative_position_biasrirDZdprrr"r#rs  F)r4r5rDZ!use_shared_relative_position_biasrrrr.Zlinspacernum_hidden_layersr ModuleListrangelayerZgradient_checkpointingrr7rr#r5s   zBeitEncoder.__init__FT)r:rroutput_hidden_statesrkr return_dictr(c Cs|rdnd}|rdnd} t|jD]\} } |r8||f}|jrv|\} } | |jj| |jjf}|j|||jdd}nd}|dur|| nd}| ||||||d}|d}|r"| |df} q"|r||f}|stdd||| fDSt||| dS) Nr"r)rkr)rrrrkrrcss|]}|dur|VqdSr3r")rvr"r"r# rz&BeitEncoder.forward..)last_hidden_stater: attentions) enumeraterrrDrLrr,rr)r6r:rrrrkrrZall_hidden_statesZall_self_attentionsrZ layer_modulerXrYrrZlayer_head_maskZ layer_outputsr"r"r#r<sB      zBeitEncoder.forward)N)NFFFNT)rrr rrrr5r.rryrwrrr<rAr"r"r7r#rs" rc@s:eZdZUeed<dZdZdZdgZdgZ dZ ddZ d S) BeitPreTrainedModelrDbeitrlTrz.*relative_position_index.*cCsLt|tjtjtjfrF|jjjd|jj d|j durB|j j nt|tj r|jjjd|jj d|j dur|jj|j  nt|tjr|j j |jjdnt|tr|jj |jdur|jj |jdur|jj nVt|tr |jj nt nt j |j |jd|_|r^t|nd|_|dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer rrN)r4r5rDrBrWrrKr}encoderuse_mean_poolingrrrrGr layernorm BeitPoolerpooler post_init)r6rDrr7r"r#r5s  zBeitModel.__init__cCs|jjSr3rWrKr=r"r"r#get_input_embeddingsszBeitModel.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)r6Zheads_to_prunerrr"r"r# _prune_headsszBeitModel._prune_headsF)rlrmrrrrkrr(c Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|||jj}|j||d\}} |jdd} |j||||| ||d} | d} | | } |j dur| | nd} |s| dur| | 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)rmr[)rrrrrrkrr)r pooler_outputr:r)rDrruse_return_dictZ get_head_maskrrWr,rrrrr:r)r6rlrmrrrrkrembedding_outputrrrZencoder_outputssequence_output pooled_outputZ head_outputsr"r"r#r<s8  zBeitModel.forward)T)NNNNFN)rrr rryr5rrrr.rrrxrrrr<rAr"r"r7r#rs( rcs8eZdZeddfdd ZejejdddZZS)rNrCcs,t|jr"tj|j|jdnd|_dS)Nr)r4r5rrrrGrrrr7r"r#r5Ss zBeitPooler.__init__r9cCsJ|jdur6|ddddddf}||d}n|dddf}|S)Nrr)rr)r6r:Z patch_tokensr r"r"r#r<Ys  zBeitPooler.forwardrr"r"r7r#rRsra 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. c s|eZdZeddfdd ZddZed eej eej eej eej ee ee e ee e e efd d d ZZS) BeitForMaskedImageModelingNrCcsTt||j|_t|dd|_tj|j|jd|_ t |j|j |_ | dS)NFrr)r4r5 num_labelsrrrrrGrrrZ vocab_sizelm_headrrr7r"r#r5ns  z#BeitForMaskedImageModeling.__init__cCsdSr3r"r=r"r"r#get_output_embeddings{sz0BeitForMaskedImageModeling.get_output_embeddingsF) rlrmrlabelsrrrkrr(c  Cs|dur |n|jj}|j|||||||d} | d} || } || ddddf} d} |durxt} | | ||} |s| f| dd}| dur| f|S|St| | | j| jdS)a bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). 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). 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)rmrrrrkrrrlosslogitsr:r) rDrrrrrrr:r)r6rlrmrrrrrkrrr Zprediction_scoresZmasked_lm_lossloss_fctr0r"r"r#r<~s4,  z"BeitForMaskedImageModeling.forward)NNNNNNFN)rrr rr5rrrr.rrxryrrrr<rAr"r"r7r#r es,   r z 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. c sleZdZeddfdd Zed eejeejeejee ee e ee e e e fdddZ ZS) BeitForImageClassificationNrCcsRt||j|_t|dd|_|jdkr 1` a classification loss is computed (Cross-Entropy). NrrrrkrrZ regressionZsingle_label_classificationZmulti_label_classificationrZr[r)rDrrrrZ problem_typer r*r.longrwr rrrgrrr:r)r6rlrrrrrkrrr rrrr0r"r"r#r<sN      "     z"BeitForImageClassification.forward)NNNNNFN)rrr rr5rrr.rryrrrr<rAr"r"r7r#rs&  rc szeZdZdZd eeeeeeeffeeeeefefeeeeeeffddfdd Z e j e j d d d Z 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. rFrN) in_channels out_channelsr{paddingrdilationr(cs<ttj||||||d|_t||_t|_dS)N)rrr{rrr) r4r5rr~conv BatchNorm2dbnZReLU activation)r6rrr{rrrr7r"r#r5's  zBeitConvModule.__init__r%r(cCs"||}||}||}|Sr3)rr!r")r6r%r0r"r"r#r<<s   zBeitConvModule.forward)rFr)rrr r!rwrrr@ryr5r.rr<rAr"r"r7r#rs rcs<eZdZeeeddfdd ZejejdddZZS)BeitPyramidPoolingBlockN) pool_scalerchannelsr(csLtt|t||ddg|_t|jD]\}}|t||q.dS)Nrr{) r4r5rZAdaptiveAvgPool2drlayersr add_moduler@)r6r%rr&rrr7r"r#r5Es   z BeitPyramidPoolingBlock.__init__r#cCs|}|jD] }||}q |Sr3)r()r6r% hidden_staterr"r"r#r<Ns  zBeitPyramidPoolingBlock.forward) rrr rwr5r.rr<rAr"r"r7r#r$Ds r$csNeZdZdZeedfeeeddfdd Zej e ej ddd 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. .N) pool_scalesrr&r_r(csht||_||_||_||_g|_t|D]2\}}t|||d}|j || t ||q0dS)N)r%rr&) r4r5r,r_rr&blocksrr$appendr)r@)r6r,rr&r_rr%blockr7r"r#r5cs  z!BeitPyramidPoolingModule.__init__)rr(cCsHg}|jD]8}||}tjj||ddd|jd}||q |S)Nr[rr\)r-rrerfr]r_r.)r6rZppm_outsppmZppm_outZupsampled_ppm_outr"r"r#r<os  z BeitPyramidPoolingModule.forward) rrr r!rrwryr5r.rlistr<rAr"r"r7r#r+Us " r+csDeZdZdZeddfdd ZddZejejdd d Z Z S) BeitUperHeadz Unified Perceptual Parsing for Scene Understanding. This head is the implementation of [UPerNet](https://huggingface.co/papers/1807.10221). Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation. NrCcst|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|qt t|j|j|jddd|_dS) NFrr'rZ)r_r r{r)r4r5r,rGrr&r_rr~r rr+ psp_modulesrr bottleneckr lateral_convs fpn_convsr.fpn_bottleneck)r6rDrZl_convZfpn_convr7r"r#r5s>    zBeitUperHead.__init__cCs:|d}|g}|||tj|dd}||}|S)NrZrr`)extendr5r.rhr6)r6inputsrZpsp_outsr0r"r"r# psp_forwards  zBeitUperHead.psp_forwardencoder_hidden_statesr(csfddtjDt}t|dddD]H}|djdd}|dtjj ||dj d|d<q@fd dt|dD}|dt|dddD]0}tjj |||djdddj d||<qt j |dd } |}|}|S) Ncsg|]\}}||qSr"r")rrZ lateral_conv)r>r"r#rrz(BeitUperHead.forward..rrrZr[rr\csg|]}j||qSr")r8r)lateralsr6r"r#rrr`)rr7r.r<rrr,rrerfr_r.rhr9r)r6r>Zused_backbone_levelsrZ prev_shapeZfpn_outsr0r")r>r?r6r#r<s$   zBeitUperHead.forward) rrr r!rr5r<r.rr<rAr"r"r7r#r2zs& r2c sTeZdZdZd eeeeeeeeffddfdd Ze j e j d d d Z Z S) BeitFCNHeada Fully Convolution Networks for Semantic Segmentation. This head is implemented of [FCNNet](https://huggingface.co/papers/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 rN)rDin_indexr{rr(c st|j|_|j|_|j|_|j|_ ||_ |d|}g}| t |j|j|||dt |jdD] }| t |j|j|||dqj|jdkrt|_n tj||_|j rt |j|j|j||dd|_tj|j|jdd|_dS)Nr[)r{rrrrr4r')r4r5rGrZauxiliary_channelsr&Zauxiliary_num_convsZ num_convsZauxiliary_concat_input concat_inputrAr.rrrrconvs Sequentialconv_catr~r r)r6rDrAr{rZ conv_paddingrCrr7r"r#r5s6     zBeitFCNHead.__init__r=cCs@||j}||}|jr2|tj||gdd}||}|S)Nrr`)rArCrBrEr.rhr)r6r>r:r0r"r"r#r<s    zBeitFCNHead.forward)r[r r) rrr r!rrwrrr5r.rr<rAr"r"r7r#r@s"r@c steZdZeddfdd ZddZed eej eej eej ee ee e ee e e e fdd d ZZS) BeitForSemanticSegmentationNrCc st||j|_t|dd|_t|jjdkr:tdt t j |j |j dddt |j t t j |j |j ddd|_t t j |j |j ddd|_t |_t jddd|_t||_|jrt|nd|_|dS)NFr r3zBeitForSemanticSegmentation 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[rz)r4r5r rrrrD out_indicesrrrDrrGr GELUfpn1fpn2rfpn3 MaxPool2dfpn4r2 decode_headZuse_auxiliary_headr@auxiliary_headrrr7r"r#r5 s*    z$BeitForSemanticSegmentation.__init__c Cstjj||jddddd}|durDtjj||jddddd}t|jjd}|||}|}|dur|||} ||jj| 7}|S)NrrFr\)Z ignore_index)rrerfr,rrDZsemantic_loss_ignore_indexZauxiliary_loss_weight) r6rauxiliary_logitsrZupsampled_logitsZupsampled_auxiliary_logitsrZ main_lossrZauxiliary_lossr"r"r# compute_loss,s  z(BeitForSemanticSegmentation.compute_lossFrcs|dur |njj}|dur |njj}|durDjjdkrDtdj|||d||d}|rd|jn|d} fddt| D} |jdjj jj fd d| D} j j j jg} tt| D]} | | | | | | <qЈ| } d}jdur| }d}|dur*| ||}|st|rJ| f|dd}n| f|d d}|durp|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))rDrG)ridxfeaturer=r"r#rvrz7BeitForSemanticSegmentation.forward..rcs<g|]4}|ddddddfddddqS)Nrrr[rZ)rdrcr)rupatch_resolutionr"r#rysr[r)rDrrr rrr:rr,rNrLrIrJrKrMrrrNrOrQrr)r6rlrrrrrkrrr>featuresopsrrrPrr0r")rurTr6r#r<?sR"       z#BeitForSemanticSegmentation.forward)NNNNNFN)rrr rr5rQrrr.rryrrrr<rAr"r"r7r#rF s(  rFzM BEiT backbone, to be used with frameworks like DETR and MaskFormer. csLeZdZfddZddZed eeeeeeee dddZ Z S) BeitBackbonec sttfddtjdD|_t|_t|jj j d|_ j rt |jjdkrrtdj}ttj||dddtj|jd ttj||ddd|_ttj||ddd|_t|_tjddd|_|dS) Ncsg|] }jqSr")rG)rrrrDr"r#rrz)BeitBackbone.__init__..rrr3zBeitBackbone 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[rzr)r4r5Z_init_backbonerrZ num_featuresrBrWrrKr}radd_fpnrrDrGrrGrrDrr Zbatch_norm_epsrHrIrJrrKrLrMr)r6rDrGr7rXr#r5s*    zBeitBackbone.__init__cCs|jjSr3rr=r"r"r#rsz!BeitBackbone.get_input_embeddingsN)rlrrrr(cCs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|jd}||\}\}}|jdd} |j|d|| |d} |r| jn| d} d} t|j | D]\\} }| |j vr|jj r|ddddddf}| ddd}| |d||}| |f7} q|jjr@|| d|| d|| d|| d g} t| } |sv|r`| f| dd}n| 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"rZr ) feature_mapsr:r)rDrrrr,rWrr:zipZ stage_namesZ out_featuresZreshape_hidden_statesrdrcrYrIrJrKrMrr r)r6rlrrrrurrsrtrrr:rZZstager*r0r"r"r#r<sP         zBeitBackbone.forward)NNN) rrr r5rrrrryr r<rAr"r"r7r#rWs rW)rr rFrrrW)r$F)Mr!collections.abcrOrrn dataclassesrtypingrrr.Ztorch.utils.checkpointrrZtorch.nnrrr Z activationsr Zmodeling_layersr Zmodeling_outputsr rrrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrZutils.backbone_utilsrZconfiguration_beitrZ get_loggerrrrr?ryr1Moduler2rBrJrrrrrrrrrrrrrr rrr$r+r2r@rFrW__all__r"r"r"r#s         f&YJ, ASL&W^N%%U;w