o ZhQ@sXdZddlZddlmZddlmZmZmZm Z m Z ddl Z ddl Z ddl m Z ddlmZmZmZddlmZdd lmZmZmZmZdd lmZmZdd lmZmZdd lm Z m!Z!m"Z"m#Z#d dl$m%Z%e"&e'Z(Gddde j)Z*Gddde j)Z+ d>de j)de j,de j,de j,dee j,de-de-fddZ.Gddde j)Z/Gdd d e j)Z0Gd!d"d"e j)Z1Gd#d$d$e j)Z2Gd%d&d&e j)Z3Gd'd(d(e j)Z4Gd)d*d*e j)Z5e!Gd+d,d,eZ6e!Gd-d.d.e6Z7Gd/d0d0e j)Z8e!d1d2Gd3d4d4e6Z9e!d5d2Gd6d7d7e6Z:eGd8d9d9e Z;e!d:d2Gd;d<dZ new_heightZ new_widthZsqrt_num_positionsr4r4r5interpolate_pos_encoding;s(     z'DeiTEmbeddings.interpolate_pos_encoding pixel_valuesbool_masked_posrFcCs|j\}}}}||}|\}} }|dur1|j|| d} |d| } |d| | | }|j|dd} |j|dd} t j | | |fdd}|j }|rW| |||}||}| |}|S)Nr9?rr=)r?r*r;r(expand unsqueezeZtype_asr&r'r$rEr,rFr/)r1rGrHrF_r7r8r6 batch_sizeZ seq_lengthZ mask_tokensmaskZ cls_tokensZdistillation_tokensZposition_embeddingr4r4r5forwardcs   zDeiTEmbeddings.forward)FNF)__name__ __module__ __qualname____doc__rboolr"r$TensorintrFr BoolTensorrO __classcell__r4r4r2r5r*s +rcs6eZdZdZfddZdejdejfddZZS)r)z 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}||_||_||_||_ t j ||||d|_ dS)Nrr) kernel_sizeZstride)r!r" image_sizer0 num_channelsr% isinstance collectionsabcIterabler+rConv2d projection)r1rr[r0r\r%r+r2r4r5r"s  zDeiTPatchEmbeddings.__init__rGrcCs<|j\}}}}||jkrtd||ddd}|S)NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.r r)r?r\ ValueErrorrbflatten transpose)r1rGrMr\r7r8xr4r4r5rOs zDeiTPatchEmbeddings.forward) rQrRrSrTr"r$rVrOrYr4r4r2r5r)s r)modulequerykeyvalueattention_maskscalingr/c Ks|t||dd|}tjj|dtjd|j}tjj |||j d}|dur,||}t||} | dd } | |fS)Nr9)r>dtype)ptrainingrr ) r$matmulrerrCZsoftmaxfloat32toror/rq contiguous) rhrirjrkrlrmr/kwargsZ attn_weightsZ attn_outputr4r4r5eager_attention_forwards  rwc sveZdZdeddffdd ZdejdejfddZ dd eejd e de e ejejfe ejffd d Z Z S)DeiTSelfAttentionrrNcst|j|jdkrt|dstd|jd|jd||_|j|_t|j|j|_|j|j|_ |j |_ |jd|_ d|_ tj|j|j |jd|_tj|j|j |jd|_tj|j|j |jd|_dS) NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .gF)bias)r!r"r%num_attention_headshasattrrcrrWattention_head_size all_head_sizeZattention_probs_dropout_prob dropout_probrm is_causalrLinearZqkv_biasrirjrkr1rr2r4r5r"s"   zDeiTSelfAttention.__init__rfcCs6|dd|j|jf}||}|ddddS)Nr9rr rr )r;r{r}rDrB)r1rfZ new_x_shaper4r4r5transpose_for_scoress z&DeiTSelfAttention.transpose_for_scoresF head_maskoutput_attentionsc Cs|||}|||}|||}t}|jjdkr4|jjdkr.|r.tdnt |jj}|||||||j |j |j sCdn|j d\}} |dd|jf} || }|rc|| f} | S|f} | S)NeagerZsdpaz`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.rg)rrmr/rn)rrjrkrirwrZ_attn_implementationloggerZ warning_oncerrrmrqrr;r~rA) r1 hidden_statesrrZ key_layerZ value_layerZ query_layerZattention_interfaceZ context_layerZattention_probsZnew_context_layer_shapeoutputsr4r4r5rOs4     zDeiTSelfAttention.forwardrP)rQrRrSrr"r$rVrrrUrrrOrYr4r4r2r5rxsrxcsFeZdZdZdeddffdd Zdejdejdejfd d ZZ S) DeiTSelfOutputz The residual connection is defined in DeiTLayer instead of here (as is the case with other models), due to the layernorm applied before each block. rrNcs.tt|j|j|_t|j|_dSN) r!r"rrr%denser-r.r/rr2r4r5r" zDeiTSelfOutput.__init__r input_tensorcC||}||}|Srrr/r1rrr4r4r5rO   zDeiTSelfOutput.forward) rQrRrSrTrr"r$rVrOrYr4r4r2r5rs$rc s~eZdZdeddffdd ZdeeddfddZ  dd ej d e ej d e de e ej ej fe ej ffd dZZS) DeiTAttentionrrNcs*tt||_t||_t|_dSr)r!r"rx attentionroutputset pruned_headsrr2r4r5r"s    zDeiTAttention.__init__headscCst|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=)lenrrr{r}rrrirjrkrrr~union)r1rindexr4r4r5 prune_headss zDeiTAttention.prune_headsFrrrcCs4||||}||d|}|f|dd}|S)Nrr)rr)r1rrrZ self_outputsattention_outputrr4r4r5rO.szDeiTAttention.forwardrP)rQrRrSrr"rrWrr$rVrrUrrrOrYr4r4r2r5rsrcs<eZdZdeddffdd ZdejdejfddZZS) DeiTIntermediaterrNcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr) r!r"rrr%intermediate_sizerr]Z hidden_actstrr intermediate_act_fnrr2r4r5r">s   zDeiTIntermediate.__init__rcCrr)rr)r1rr4r4r5rOFrzDeiTIntermediate.forward rQrRrSrr"r$rVrOrYr4r4r2r5r=srcsBeZdZdeddffdd Zdejdejdejfdd ZZS) DeiTOutputrrNcs.tt|j|j|_t|j|_ dSr) r!r"rrrr%rr-r.r/rr2r4r5r"OrzDeiTOutput.__init__rrcCs ||}||}||}|Srrrr4r4r5rOTs  zDeiTOutput.forwardrr4r4r2r5rNs$rc sleZdZdZdeddffdd Z  d dejd eejd e de e ejejfe ejffd d Z Z S) DeiTLayerz?This corresponds to the Block class in the timm implementation.rrNcsbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dS)NrZeps)r!r"Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr LayerNormr%layer_norm_epslayernorm_beforelayernorm_afterrr2r4r5r"as    zDeiTLayer.__init__FrrrcCs`|j||||d}|d}|dd}||}||}||}|||}|f|}|S)N)rrr)rrrrr)r1rrrZself_attention_outputsrrZ layer_outputr4r4r5rOks     zDeiTLayer.forwardrP)rQrRrSrTrr"r$rVrrUrrrOrYr4r4r2r5r^s rcsbeZdZdeddffdd Z    ddejd eejd ed ed ede e e ff d dZ Z S) DeiTEncoderrrNcs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr4)r).0rLrr4r5 sz(DeiTEncoder.__init__..F) r!r"rrZ ModuleListrangenum_hidden_layerslayergradient_checkpointingrr2rr5r"s   zDeiTEncoder.__init__FTrrroutput_hidden_states return_dictc Cs|rdnd}|r dnd}t|jD]8\}} |r||f}|dur$||nd} |jr6|jr6|| j|| |} n| || |} | d}|rI|| df}q|rQ||f}|s_tdd|||fDSt|||dS)Nr4rrcss|] }|dur|VqdSrr4)rvr4r4r5 sz&DeiTEncoder.forward..)last_hidden_stater attentions) enumeraterrrqZ_gradient_checkpointing_func__call__tupler) r1rrrrrZall_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskZ layer_outputsr4r4r5rOs6      zDeiTEncoder.forward)NFFT)rQrRrSrr"r$rVrrUrrrrOrYr4r4r2r5rs&  rc@sLeZdZeZdZdZdZdgZdZ dZ de e j e je jfddfdd ZdS) DeiTPreTrainedModeldeitrGTrrhrNcCst|tjtjfr0tjj|jjt j d|j j d|jj |j_|jdur.|jjdSdSt|tjrE|jj|jjddSt|tri|jj|jj|jj|jdurk|jjdSdSdS)zInitialize the weightsrg)meanZstdNrI)r]rrrainitZ trunc_normal_weightdatartr$rsrZinitializer_rangerorzZzero_rZfill_rr&r,r'r()r1rhr4r4r5 _init_weightss(         z!DeiTPreTrainedModel._init_weights)rQrRrSrZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesZ_supports_sdpaZ_supports_flash_attn_2rrrrarrr4r4r4r5rs&rcseZdZddedededdffdd Zdefd d Zd d Ze       dde e j de e j de e j de ede ede ededeeeffddZZS) DeiTModelTFradd_pooling_layerrrNcs\t|||_t||d|_t||_tj|j |j d|_ |r%t |nd|_ |dS)z add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer use_mask_token (`bool`, *optional*, defaults to `False`): Whether to use a mask token for masked image modeling. )rrN)r!r"rrr6rencoderrrr%r layernorm DeiTPoolerpooler post_init)r1rrrr2r4r5r"s   zDeiTModel.__init__cCs|jjSr)r6r*)r1r4r4r5get_input_embeddingsszDeiTModel.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)r1Zheads_to_prunerrr4r4r5 _prune_headsszDeiTModel._prune_headsrGrHrrrrrFcCs |dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|||jj}|jjj j j }|j |kr?| |}|j|||d} |j | ||||d} | d} || } |jdurd|| nd} |s{| durp| | fn| f} | | ddSt| | | j| jdS)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). Nz You have to specify pixel_values)rHrF)rrrrrr)rZ pooler_outputrr)rrruse_return_dictrcZ get_head_maskrr6r*rbrrortrrrrrr)r1rGrHrrrrrFZexpected_dtypeZembedding_outputZencoder_outputssequence_output pooled_outputZ head_outputsr4r4r5rOs@   zDeiTModel.forward)TFNNNNNNF)rQrRrSrrUr"r)rrrrr$rVrXrrrrOrYr4r4r2r5rs:  rcs*eZdZdeffdd ZddZZS)rrcs,tt|j|j|_t|j|_ dSr) r!r"rrr%Zpooler_output_sizerr Z pooler_act activationrr2r4r5r"<s zDeiTPooler.__init__cCs(|dddf}||}||}|S)Nr)rr)r1rZfirst_token_tensorrr4r4r5rOAs  zDeiTPooler.forward)rQrRrSrr"rOrYr4r4r2r5r;sra\ DeiT Model with a decoder on top for masked image modeling, as proposed in [SimMIM](https://arxiv.org/abs/2111.09886). Note that we provide a script to pre-train this model on custom data in our [examples directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining). )Z custom_introcseZdZdeddffdd Ze       ddeejdeej d eejd ee d ee d ee d e de e e ffddZZS)DeiTForMaskedImageModelingrrNcsXt|t|ddd|_ttj|j|jd|j ddt |j|_ | dS)NFT)rrr r)Z in_channelsZ out_channelsrZ) r!r"rrrZ Sequentialrar%Zencoder_strider\Z PixelShuffledecoderrrr2r4r5r"Ws   z#DeiTForMaskedImageModeling.__init__FrGrHrrrrrFc CsJ|dur|n|jj}|j|||||||d}|d} | ddddf} | j\} } } t| d} }| ddd| | | |} || }d}|dur|jj|jj }|d||}| |jj d |jj d d }t jj||dd }|||d |jj}|s|f|dd}|dur|f|S|St|||j|jd S) 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). Examples: ```python >>> from transformers import AutoImageProcessor, DeiTForMaskedImageModeling >>> 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("facebook/deit-base-distilled-patch16-224") >>> model = DeiTForMaskedImageModeling.from_pretrained("facebook/deit-base-distilled-patch16-224") >>> 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, reconstructed_pixel_values = outputs.loss, outputs.reconstruction >>> list(reconstructed_pixel_values.shape) [1, 3, 224, 224] ```N)rHrrrrrFrrr9r:r none)Z reductiongh㈵>)lossZreconstructionrr)rrrr?rWrBrArr[r0Zrepeat_interleaverKrurrCZl1_losssumr\rrr)r1rGrHrrrrrFrrrMZsequence_lengthr\r7r8Zreconstructed_pixel_valuesZmasked_im_lossr;rNZreconstruction_lossrr4r4r5rOhsH&     z"DeiTForMaskedImageModeling.forwardr)rQrRrSrr"rrr$rVrXrUrrrrOrYr4r4r2r5rJs6  rz DeiT Model transformer with an image classification head on top (a linear layer on top of the final hidden state of the [CLS] token) e.g. for ImageNet. cseZdZdeddffdd Ze       ddeejdeejd eejd ee d ee d ee d e de e e ffddZ ZS)DeiTForImageClassificationrrNcsRt||j|_t|dd|_|jdkrt|j|jnt|_ | dSNF)rr) r!r" num_labelsrrrrr%Identity classifierrrr2r4r5r"s $ z#DeiTForImageClassification.__init__FrGrlabelsrrrrFcCs|dur|n|jj}|j||||||d}|d} || dddddf} d} |dur|| j}|jjdur]|jdkrCd|j_n|jdkrY|jt j ksT|jt j krYd|j_nd|j_|jjdkr{t } |jdkru| | | } n+| | |} n%|jjdkrt} | | d|j|d} n|jjdkrt} | | |} |s| f|dd} | dur| f| S| St| | |j|jd S) aZ 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, DeiTForImageClassification >>> import torch >>> from PIL import Image >>> import requests >>> torch.manual_seed(3) # doctest: +IGNORE_RESULT >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> # note: we are loading a DeiTForImageClassificationWithTeacher from the hub here, >>> # so the head will be randomly initialized, hence the predictions will be random >>> image_processor = AutoImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224") >>> model = DeiTForImageClassification.from_pretrained("facebook/deit-base-distilled-patch16-224") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> # model predicts one of the 1000 ImageNet classes >>> predicted_class_idx = logits.argmax(-1).item() >>> print("Predicted class:", model.config.id2label[predicted_class_idx]) Predicted class: Polaroid camera, Polaroid Land camera ```NrrrrrFrrZ regressionZsingle_label_classificationZmulti_label_classificationr9)rlogitsrr)rrrrrtZdeviceZ problem_typerror$longrWr Zsqueezer rDr rrr)r1rGrrrrrrFrrrrZloss_fctrr4r4r5rOsP*     "       z"DeiTForImageClassification.forwardr)rQrRrSrr"rrr$rVrUrrrrOrYr4r4r2r5rs6  rc@steZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)+DeiTForImageClassificationWithTeacherOutputa5 Output type of [`DeiTForImageClassificationWithTeacher`]. Args: logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Prediction scores as the average of the cls_logits and distillation logits. cls_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the class token). distillation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the distillation token). 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. Nr cls_logitsdistillation_logitsrr)rQrRrSrTrrr$Z FloatTensor__annotations__rrrrrr4r4r4r5r.s ra DeiT Model transformer with image classification heads on top (a linear layer on top of the final hidden state of the [CLS] token and a linear layer on top of the final hidden state of the distillation token) e.g. for ImageNet. .. warning:: This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet supported. cs~eZdZdeddffdd Ze      ddeejdeejd ee d ee d ee d e de e e ffd dZ ZS)%DeiTForImageClassificationWithTeacherrrNcsvt||j|_t|dd|_|jdkrt|j|jnt|_ |jdkr0t|j|jnt|_ | dSr) r!r"rrrrrr%rcls_classifierdistillation_classifierrrr2r4r5r"Ys    z.DeiTForImageClassificationWithTeacher.__init__FrGrrrrrFc Cs|dur|n|jj}|j||||||d}|d}||dddddf} ||dddddf} | | d} |sJ| | | f|dd} | St| | | |j|jdS)Nrrrr )rrrrr)rrrrrrrr) r1rGrrrrrFrrrrrrr4r4r5rOjs.   z-DeiTForImageClassificationWithTeacher.forward)NNNNNF)rQrRrSrr"rrr$rVrUrrrrOrYr4r4r2r5rMs0  r)rrrrr)rg)>rTcollections.abcr^ dataclassesrtypingrrrrrr$Ztorch.utils.checkpointrZtorch.nnr r r Z activationsr Zmodeling_outputsrrrrZmodeling_utilsrrZ pytorch_utilsrrutilsrrrrZconfiguration_deitrZ get_loggerrQrModulerr)rVfloatrwrxrrrrrrrrrrrrr__all__r4r4r4r5sx     Y( ?(+3_ hj <