o hJ@sHdZddlZddlZddlmZddlmZmZm Z m Z ddl Z ddl Z ddl m Z ddlmZmZmZddlmZdd lmZmZmZmZdd lmZdd lmZmZdd lmZm Z m!Z!m"Z"m#Z#m$Z$d dl%m&Z&e#'e(Z)dZ*dZ+gdZ,dZ-dZ.dgZ/Gddde j0Z1Gddde j0Z2Gddde j0Z3Gddde j0Z4Gddde j0Z5Gddde j0Z6Gdd d e j0Z7Gd!d"d"e j0Z8Gd#d$d$e j0Z9Gd%d&d&eZ:d'Z;d(Ze!d.e;Gd/d0d0e:Z?e!d1e;Gd2d3d3e:Z@eGd4d5d5eZAe!d6e;Gd7d8d8e:ZBdS)9z PyTorch DeiT model.N) dataclass)OptionalSetTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BaseModelOutputBaseModelOutputWithPoolingImageClassifierOutputMaskedImageModelingOutput)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) DeiTConfigrz(facebook/deit-base-distilled-patch16-224)riztabby, tabby catcsReZdZdZd dededdffdd Zdd ejd e ej dejfd d Z Z S)DeiTEmbeddingszv Construct the CLS token, distillation token, position and patch embeddings. Optionally, also the mask token. Fconfiguse_mask_tokenreturnNcstttdd|j|_ttdd|j|_|r*ttdd|jnd|_ t ||_ |j j }ttd|d|j|_ t|j|_dS)Nr)super__init__r Parametertorchzeros hidden_size cls_tokendistillation_token mask_tokenDeiTPatchEmbeddingspatch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropout)selfrrr- __class__\/var/www/html/ai/venv/lib/python3.10/site-packages/transformers/models/deit/modeling_deit.pyr#Is   zDeiTEmbeddings.__init__ pixel_valuesbool_masked_posc Cs||}|\}}}|dur*|j||d}|d|}|d|||}|j|dd} |j|dd} tj | | |fdd}||j }| |}|S)N?rdim) r,sizer*expand unsqueezetype_asr(r)r%catr.r1) r2r7r8 embeddings batch_size seq_length_ mask_tokensmask cls_tokensdistillation_tokensr5r5r6forwardTs   zDeiTEmbeddings.forward)FN) __name__ __module__ __qualname____doc__rboolr#r%Tensorr BoolTensorrJ __classcell__r5r5r3r6rDs* 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_sizestride)r"r# image_size patch_size num_channelsr' isinstance collectionsabcIterabler-rConv2d projection)r2rrVrWrXr'r-r3r5r6r#ms  zDeiTPatchEmbeddings.__init__r7r c Cs|j\}}}}||jkrtd||jdks||jdkr5td|d|d|jdd|jdd ||ddd}|S) NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rrzInput image size (*z) doesn't match model (z).r!)shaperX ValueErrorrVr^flatten transpose)r2r7rCrXheightwidthxr5r5r6rJ|s (zDeiTPatchEmbeddings.forward) rLrMrNrOr#r%rQrJrSr5r5r3r6r+fs r+c sveZdZdeddffdd ZdejdejfddZ dd eejd e de e ejejfe ejffd d Z Z S)DeiTSelfAttentionrr Ncst|j|jdkr t|ds td|jfd|jd|j|_t|j|j|_|j|j|_t j |j|j|j d|_ t j |j|j|j d|_ t j |j|j|j d|_t |j|_dS)Nrembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .)bias)r"r#r'num_attention_headshasattrraintattention_head_size all_head_sizerLinearqkv_biasquerykeyvaluer/attention_probs_dropout_probr1r2rr3r5r6r#s  zDeiTSelfAttention.__init__rfcCs6|dd|j|jf}||}|ddddS)Nr9rr!rr )r=rkrnviewpermute)r2rf new_x_shaper5r5r6transpose_for_scoress z&DeiTSelfAttention.transpose_for_scoresF head_maskoutput_attentionsc Cs||}|||}|||}||}t||dd}|t|j }t j j |dd} | | } |durA| |} t| |} | dddd} | dd|jf} | | } |rj| | f} | S| f} | S)Nr9r;rr!rr )rrrzrsrtr%matmulrcmathsqrtrnr functionalsoftmaxr1rx contiguousr=rorw) r2 hidden_statesr{r|mixed_query_layer key_layer value_layer query_layerattention_scoresattention_probs context_layernew_context_layer_shapeoutputsr5r5r6rJs$      zDeiTSelfAttention.forwardNF)rLrMrNrr#r%rQrzrrPrrrJrSr5r5r3r6rgsrgcsFeZdZdZdeddffdd 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. rr Ncs.tt|j|j|_t|j|_dSrK) r"r#rrpr'denser/r0r1rvr3r5r6r# zDeiTSelfOutput.__init__r input_tensorcC||}||}|SrKrr1r2rrr5r5r6rJ  zDeiTSelfOutput.forward) rLrMrNrOrr#r%rQrJrSr5r5r3r6rs$rc s~eZdZdeddffdd ZdeeddfddZ  dd ej d e ej d e de e ej ej fe ej ffd dZZS) DeiTAttentionrr Ncs*tt||_t||_t|_dSrK)r"r#rg attentionroutputset pruned_headsrvr3r5r6r#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;)lenrrrkrnrrrrrsrtrrrounion)r2rindexr5r5r6 prune_headss zDeiTAttention.prune_headsFrr{r|cCs4||||}||d|}|f|dd}|S)Nrr)rr)r2rr{r| self_outputsattention_outputrr5r5r6rJszDeiTAttention.forwardr)rLrMrNrr#rrmrr%rQrrPrrrJrSr5r5r3r6rsrcs<eZdZdeddffdd ZdejdejfddZZS) DeiTIntermediaterr NcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSrK) r"r#rrpr'intermediate_sizerrY hidden_actstrr intermediate_act_fnrvr3r5r6r#s   zDeiTIntermediate.__init__rcCrrK)rr)r2rr5r5r6rJ rzDeiTIntermediate.forward rLrMrNrr#r%rQrJrSr5r5r3r6rsrcsBeZdZdeddffdd Zdejdejdejfdd ZZS) DeiTOutputrr Ncs.tt|j|j|_t|j|_ dSrK) r"r#rrprr'rr/r0r1rvr3r5r6r#rzDeiTOutput.__init__rrcCs ||}||}||}|SrKrrr5r5r6rJs  zDeiTOutput.forwardrr5r5r3r6rs$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.rr Ncsbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dS)Nreps)r"r#chunk_size_feed_forward seq_len_dimrrr intermediaterrr LayerNormr'layer_norm_epslayernorm_beforelayernorm_afterrvr3r5r6r#'s    zDeiTLayer.__init__Frr{r|cCs`|j||||d}|d}|dd}||}||}||}|||}|f|}|S)N)r|rr)rrrrr)r2rr{r|self_attention_outputsrr layer_outputr5r5r6rJ1s     zDeiTLayer.forwardr)rLrMrNrOrr#r%rQrrPrrrJrSr5r5r3r6r$s rcsbeZdZdeddffdd Z    ddejd eejd ed ed ede e e ff d dZ Z S) DeiTEncoderrr Ncs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr5)r).0rErr5r6 Ssz(DeiTEncoder.__init__..F) r"r#rr ModuleListrangenum_hidden_layerslayergradient_checkpointingrvr3rr6r#Ps   zDeiTEncoder.__init__FTrr{r|output_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)Nr5rrcss|] }|dur|VqdSrKr5)rvr5r5r6 zsz&DeiTEncoder.forward..)last_hidden_stater attentions) enumeraterrtraining_gradient_checkpointing_func__call__tupler ) r2rr{r|rrall_hidden_statesall_self_attentionsi layer_modulelayer_head_mask layer_outputsr5r5r6rJVs6      zDeiTEncoder.forward)NFFT)rLrMrNrr#r%rQrrPrrr rJrSr5r5r3r6rOs&  rc@sHeZdZdZeZdZdZdZdgZ de e j e j e jfddfd d ZdS) DeiTPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. deitr7Trmoduler NcCst|tjtjfr0tjj|jjt j d|j j d|jj |j_|jdur.|jjdSdSt|tjrE|jj|jjddSdS)zInitialize the weightsg)meanstdNr:)rYrrpr]init trunc_normal_weightdatator%float32rinitializer_rangedtyperjzero_rfill_)r2rr5r5r6 _init_weightss    z!DeiTPreTrainedModel._init_weights)rLrMrNrOr config_classbase_model_prefixmain_input_namesupports_gradient_checkpointing_no_split_modulesrrrpr]rrr5r5r5r6rs&raF This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`DeiTConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. aL Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`DeiTImageProcessor.__call__`] for details. head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. z^The bare DeiT Model transformer outputting raw hidden-states without any specific head on top.cseZdZddedededdffdd Zdefd d Zd d Ze e e e e eded      ddeejdeejdeejdeedeedeedeee ffddZZS) DeiTModelTFradd_pooling_layerrr Ncs\t|||_t||d|_t||_tj|j |j d|_ |r%t |nd|_ |dS)N)rr)r"r#rrrBrencoderrrr'r layernorm DeiTPoolerpooler post_init)r2rrrr3r5r6r#s   zDeiTModel.__init__cCs|jjSrK)rBr,)r2r5r5r6get_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)r2heads_to_prunerrr5r5r6 _prune_headsszDeiTModel._prune_headsvision) checkpoint output_typermodalityexpected_outputr7r8r{r|rrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|||jj}|jjj j j }|j |kr?| |}|j||d}|j |||||d} | d} || } |jdurc|| nd} |sz| duro| | 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)r8r{r|rrrr)r pooler_outputrr)rr|ruse_return_dictra get_head_maskrrBr,r^rrrrrrrrr) r2r7r8r{r|rrexpected_dtypeembedding_outputencoder_outputssequence_output pooled_output head_outputsr5r5r6rJs<   zDeiTModel.forward)TFNNNNNN)rLrMrNrrPr#r+rrrDEIT_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPErr%rQrRrrrJrSr5r5r3r6rsB    rcs*eZdZdeffdd ZddZZS)rrcs*tt|j|j|_t|_dSrK)r"r#rrpr'rTanh activationrvr3r5r6r#!s zDeiTPooler.__init__cCs(|dddf}||}||}|S)Nr)rr)r2rfirst_token_tensorrr5r5r6rJ&s  zDeiTPooler.forward)rLrMrNrr#rJrSr5r5r3r6r sraWDeiT 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). cseZdZdeddffdd Zeeeee d      dde e j de e j d e e j d e ed e ed e edeeeffd dZZS)DeiTForMaskedImageModelingrr NcsXt|t|ddd|_ttj|j|jd|j ddt |j|_ | dS)NFT)rrr!r) in_channels out_channelsrT) r"r#rrr Sequentialr]r'encoder_striderX PixelShuffledecoderrrvr3r5r6r#<s   z#DeiTForMaskedImageModeling.__init__rrr7r8r{r|rrcCsH|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) aM 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). Returns: 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)r8r{r|rrrrr9g?r!none) reductiongh㈵>)lossreconstructionrr)rrrr`rmrxreshaper rVrWrepeat_interleaver?rrrl1_losssumrXrrr)r2r7r8r{r|rrrrrCsequence_lengthrXrdrereconstructed_pixel_valuesmasked_im_lossr=rGreconstruction_lossrr5r5r6rJMsF(     z"DeiTForMaskedImageModeling.forwardr)rLrMrNrr#rrrrrrr%rQrRrPrrrJrSr5r5r3r6r/s2   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 Zeeeee d      dde e j de e j d e e j d e e d e e d e e deeeffd dZZS)DeiTForImageClassificationrr NcsRt||j|_t|dd|_|jdkrt|j|jnt|_ | dSNF)rr) r"r# num_labelsrrrrpr'Identity classifierrrvr3r5r6r#s $ z#DeiTForImageClassification.__init__rr7r{labelsr|rrc Cs~|dur|n|jj}|j|||||d}|d}||dddddf} d} |dur|| j}|jjdur\|jdkrBd|j_n|jdkrX|jt j ksS|jt j krXd|j_nd|j_|jjdkrzt } |jdkrt| | | } n+| | |} n%|jjdkrt} | | d|j|d} n|jjdkrt} | | |} |s| f|dd} | dur| f| S| St| | |j|jd S) aM labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the image classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). Returns: 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: magpie ```Nrrr regressionsingle_label_classificationmulti_label_classificationr9)rlogitsrr)rrrrrdevice problem_typerrr%longrmr squeezer rwrrrr) r2r7r{r r|rrrrr$rloss_fctrr5r5r6rJsN,    "       z"DeiTForImageClassification.forwardr)rLrMrNrr#rrrrrrr%rQrPrrrJrSr5r5r3r6rs2  rc@sheZdZUdZdZejed<dZejed<dZ ejed<dZ e e ejed<dZ e e 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)rLrMrNrOr$r% FloatTensor__annotations__r+r,rrrrr5r5r5r6r*s r*a 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. cseZdZdeddffdd Zeeeee e e d     dde e jde e jd e ed e ed e edeee ff d d ZZS)%DeiTForImageClassificationWithTeacherrr Ncsvt||j|_t|dd|_|jdkrt|j|jnt|_ |jdkr0t|j|jnt|_ | dSr) r"r#rrrrrpr'rcls_classifierdistillation_classifierrrvr3r5r6r#Bs    z.DeiTForImageClassificationWithTeacher.__init__)rrrrr7r{r|rrc Cs|dur|n|jj}|j|||||d}|d}||dddddf}||dddddf} || d} |sI| || f|dd} | St| || |j|jdS)Nrrrr!)r$r+r,rr)rrrr0r1r*rr) r2r7r{r|rrrrr+r,r$rr5r5r6rJSs, z-DeiTForImageClassificationWithTeacher.forward)NNNNN)rLrMrNrr#rrr_IMAGE_CLASS_CHECKPOINTr*r_IMAGE_CLASS_EXPECTED_OUTPUTrr%rQrPrrrJrSr5r5r3r6r/5s6  r/)CrOcollections.abcrZr dataclassesrtypingrrrrr%torch.utils.checkpointrtorch.nnrr r activationsr modeling_outputsr rrrmodeling_utilsr pytorch_utilsrrutilsrrrrrrconfiguration_deitr get_loggerrLloggerrrrr2r3"DEIT_PRETRAINED_MODEL_ARCHIVE_LISTModulerr+rgrrrrrrrDEIT_START_DOCSTRINGrrrrrr*r/r5r5r5r6st       "%=(+3 ]  ik