U 9%eR@sJdZddlmZmZmZddlZddlZddlmZddlm Z m Z m Z ddl m Z ddlmZmZmZmZdd lmZdd lmZmZmZmZmZd d lmZeeZd Z dZ!d dddgZ"dZ#dZ$dgZ%d?e&e&ee&e&dddZ'e(de(dfe(e(e(e(dddZ)Gdddej*Z+Gdddej*Z,Gdddej*Z-Gd d!d!ej*Z.Gd"d#d#ej*Z/Gd$d%d%ej*Z0Gd&d'd'ej*Z1Gd(d)d)ej*Z2Gd*d+d+ej*Z3Gd,d-d-eZ4d.Z5d/Z6ed0e5Gd1d2d2e4Z7ed3e5Gd4d5d5e4Z8Gd6d7d7ej*Z9Gd8d9d9ej*Z:Gd:d;d;ej*Z;edd>e4Z||7}t|S)a Ensure that all layers have a channel count that is divisible by `divisor`. This function is taken from the original TensorFlow repo. It can be seen here: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py Ng?)maxint)rrr new_valuer!s/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/mobilevitv2/modeling_mobilevitv2.pymake_divisibleCs  r#z-infinf)rmin_valmax_valrcCst|t||SN)rminrr%r&r!r!r"clipRsr*c sTeZdZd eeeeeeeeeeeefdd fdd Ze j e j dd d Z Z S) MobileViTV2ConvLayerrFTN) config in_channels out_channels kernel_sizestridegroupsbiasdilationuse_normalizationuse_activationrc stt|dd|} ||dkr@td|d|d||dkrbtd|d|dtj||||| |||dd |_| rtj|d d d d d |_nd|_| rt | t rt | |_ qt |j t rt |j |_ q|j |_ nd|_ dS)NrrrzInput channels (z) are not divisible by z groups.zOutput channels (Zzeros) r-r.r/r0paddingr3r1r2Z padding_modegh㈵>g?T)Z num_featuresepsZmomentumZaffineZtrack_running_stats)super__init__r ValueErrorrConv2d convolutionZ BatchNorm2d normalization isinstancestrr activationZ hidden_act) selfr,r-r.r/r0r1r2r3r4r5r6 __class__r!r"r9XsB         zMobileViTV2ConvLayer.__init__featuresrcCs6||}|jdk r||}|jdk r2||}|Sr')r<r=r@)rArEr!r!r"forwards      zMobileViTV2ConvLayer.forward)rrFrTT) __name__ __module__ __qualname__rrboolrr?r9torchTensorrF __classcell__r!r!rBr"r+Ws( 6r+csFeZdZdZd eeeeeddfdd Zejejddd Z Z S) MobileViTV2InvertedResidualzQ Inverted residual block (MobileNetv2): https://arxiv.org/abs/1801.04381 rN)r,r-r.r0r3rc stttt||jd}|dkr:td|d|dkoH||k|_t|||dd|_ t|||d|||d|_ t|||dd d |_ dS) Nr)rrzInvalid stride .r)r-r.r/r )r-r.r/r0r1r3Fr-r.r/r5) r8r9r#rroundZ expand_ratior: use_residualr+ expand_1x1conv_3x3 reduce_1x1)rAr,r-r.r0r3Zexpanded_channelsrBr!r"r9s6  z$MobileViTV2InvertedResidual.__init__rDcCs4|}||}||}||}|jr0||S|Sr')rSrTrUrR)rArEZresidualr!r!r"rFs    z#MobileViTV2InvertedResidual.forward)r rGrHrI__doc__rrr9rKrLrFrMr!r!rBr"rNs!rNcsBeZdZd eeeeeddfdd ZejejdddZZ S) MobileViTV2MobileNetLayerrN)r,r-r.r0 num_stagesrcsRtt|_t|D]0}t||||dkr4|ndd}|j||}qdS)Nrr)r-r.r0)r8r9r ModuleListlayerrangerNappend)rAr,r-r.r0rYir[rBr!r"r9s    z"MobileViTV2MobileNetLayer.__init__rDcCs|jD] }||}q|Sr'r[)rArE layer_moduler!r!r"rFs  z!MobileViTV2MobileNetLayer.forward)rr rGrHrIrrr9rKrLrFrMr!r!rBr"rXsrXcs>eZdZdZeeddfdd ZejejdddZ Z S) MobileViTV2LinearSelfAttentionaq This layer applies a self-attention with linear complexity, as described in MobileViTV2 paper: https://arxiv.org/abs/2206.02680 Args: config (`MobileVitv2Config`): Model configuration object embed_dim (`int`): `input_channels` from an expected input of size :math:`(batch_size, input_channels, height, width)` N)r, embed_dimrc s\tt||dd|ddddd|_tj|jd|_t|||ddddd|_||_dS)NrrTF)r,r-r.r2r/r4r5p) r8r9r+qkv_projrDropout attn_dropoutout_projrc)rAr,rcrBr!r"r9s*    z'MobileViTV2LinearSelfAttention.__init__ hidden_statesrc Cs||}tj|d|j|jgdd\}}}tjjj|dd}||}||}tj|ddd}tjj || |}| |}|S)Nr)Zsplit_size_or_sectionsdimrlTrlZkeepdim) rfrKsplitrcr functionalZsoftmaxrhsumreluZ expand_asri) rArkZqkvquerykeyrZcontext_scoresZcontext_vectoroutr!r!r"rFs    z&MobileViTV2LinearSelfAttention.forwardrVr!r!rBr"rbs rbcs@eZdZd eeeeddfdd ZejejdddZ Z S) MobileViTV2FFNN)r,rcffn_latent_dim ffn_dropoutrc sZtt|||dddddd|_t||_t|||dddddd|_t||_dS)NrTF)r,r-r.r/r0r2r4r5) r8r9r+conv1rrgdropout1conv2dropout2)rAr,rcryrzrBr!r"r9 s.   zMobileViTV2FFN.__init__rjcCs,||}||}||}||}|Sr')r{r|r}r~)rArkr!r!r"rF@s     zMobileViTV2FFN.forward)rx rGrHrIrrfloatr9rKrLrFrMr!r!rBr"rws rwcs@eZdZd eeeeddfdd ZejejdddZ Z S) MobileViTV2TransformerLayerrxN)r,rcrydropoutrcsbttjd||jd|_t|||_tj|d|_ tjd||jd|_ t ||||j |_ dS)NrZ num_groups num_channelsr7rd)r8r9r GroupNormlayer_norm_epslayernorm_beforerb attentionrgr|layernorm_afterrwrzffn)rAr,rcryrrBr!r"r9Is   z$MobileViTV2TransformerLayer.__init__rjcCs<||}||}||}||}||}||}|Sr')rrrr)rArkZlayernorm_1_outZattention_outputZ layer_outputr!r!r"rFWs    z#MobileViTV2TransformerLayer.forward)rxrr!r!rBr"rHsrcs<eZdZeeeddfdd ZejejdddZZ S)MobileViTV2TransformerN)r,n_layersd_modelrcsft|j}||g|}dd|D}t|_t|D]"}t||||d}|j|q>dS)NcSsg|]}t|ddqS))r).0dr!r!r" lsz3MobileViTV2Transformer.__init__..)rcry) r8r9ffn_multiplierrrZr[r\rr])rAr,rrrZffn_dimsZ block_idxtransformer_layerrBr!r"r9ds   zMobileViTV2Transformer.__init__rjcCs|jD] }||}q|Sr'r_)rArkr`r!r!r"rFus  zMobileViTV2Transformer.forwardrar!r!rBr"rcsrc seZdZdZdeeeeeeeddfdd Zeje eje eeffdd d Z eje eefejd d d Z ejejdddZ Z S)MobileViTV2Layerz= MobileViTV2 layer: https://arxiv.org/abs/2206.02680 rrN)r,r-r. attn_unit_dim n_attn_blocksr3r0rc st|j|_|j|_|}|dkr\t||||dkr:|nd|dkrL|dndd|_|}nd|_t||||j|d|_ t|||dddd|_ t |||d|_ t jd||jd|_t|||dd dd|_dS) Nrr)r-r.r0r3)r-r.r/r1F)r-r.r/r4r5)rrrT)r8r9 patch_size patch_width patch_heightrNdownsampling_layerr+Zconv_kernel_sizeconv_kxkconv_1x1r transformerrrr layernormconv_projection) rAr,r-r.rrr3r0Z cnn_out_dimrBr!r"r9sN  zMobileViTV2Layer.__init__) feature_maprcCsT|j\}}}}tjj||j|jf|j|jfd}||||j|jd}|||ffS)N)r/r0rm)shaperrqZunfoldrrreshape)rAr batch_sizer-Z img_heightZ img_widthpatchesr!r!r" unfoldings  zMobileViTV2Layer.unfolding)r output_sizercCsH|j\}}}}|||||}tjj|||j|jf|j|jfd}|S)N)rr/r0)rrrrqfoldrr)rArrrZin_dimrZ n_patchesrr!r!r"foldings  zMobileViTV2Layer.foldingrDcCs`|jr||}||}||}||\}}||}||}|||}||}|Sr')rrrrrrrr)rArErrr!r!r"rFs       zMobileViTV2Layer.forward)rrr)rGrHrIrWrrr9rKrLrrrrFrMr!r!rBr"r{s" =$  rcsDeZdZeddfdd Zd ejeeee e fddd Z Z S) MobileViTV2EncoderNr,rcst||_t|_d|_d}}|jdkr.create_custom_forward..custom_forwardr!)rrr!rr"create_custom_forwardJs z9MobileViTV2Encoder.forward..create_custom_forwardcss|]}|dk r|VqdSr'r!)rvr!r!r" [sz-MobileViTV2Encoder.forward..)last_hidden_staterk) enumerater[rZtrainingrKutils checkpointtupler )rArkrrZall_hidden_statesr^r`rr!r!r"rF?s   zMobileViTV2Encoder.forward)FT) rGrHrIrr9rKrLrJrrr rFrMr!r!rBr"rsT rc@sJeZdZdZeZdZdZdZe e j e j e j fddddZd d d ZdS) MobileViTV2PreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. mobilevitv2 pixel_valuesTN)rrcCsjt|tjtjfr@|jjjd|jjd|j dk rf|j j n&t|tj rf|j j |jj ddS)zInitialize the weightsrx)meanZstdNg?) r>rLinearr;weightdataZnormal_r,Zinitializer_ranger2Zzero_ LayerNormZfill_)rArr!r!r" _init_weightsls   z(MobileViTV2PreTrainedModel._init_weightsFcCst|tr||_dSr')r>rr)rArrr!r!r"_set_gradient_checkpointingxs z6MobileViTV2PreTrainedModel._set_gradient_checkpointing)F)rGrHrIrWr config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrrrr;rrrr!r!r!r"ras raM 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 ([`MobileViTV2Config`]): 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. aF Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`MobileViTImageProcessor.__call__`] for details. 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. zYThe bare MobileViTV2 model outputting raw hidden-states without any specific head on top.c steZdZd eedfdd ZddZeee e e e de ddeejeeeeeee fd d d ZZS)MobileViTV2ModelT)r, expand_outputc sft|||_||_ttd|jdddddd}t||j|ddd d d |_ t ||_ | dS) Nrrrr)rrr rTr-r.r/r0r4r5) r8r9r,rr#r*rr+r conv_stemrencoder post_init)rAr,rrrBr!r"r9s&  zMobileViTV2Model.__init__cCsF|D]8\}}|jj|}t|tr|jjD]}|j|q.qdS)zPrunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsrr[r>rrrZ prune_heads)rAZheads_to_pruneZ layer_indexZheadsZmobilevitv2_layerrr!r!r" _prune_headss    zMobileViTV2Model._prune_headsZvision)r output_typerZmodalityexpected_outputN)rrrrc Cs|dk r |n|jj}|dk r |n|jj}|dkr8td||}|j|||d}|jrv|d}tj|ddgdd}n |d}d}|s|dk r||fn|f}||ddSt |||j d S) Nz You have to specify pixel_valuesrrrrmFror)r pooler_outputrk) r,ruse_return_dictr:rrrrKrr rk) rArrrZembedding_outputZencoder_outputsr pooled_outputoutputr!r!r"rFs0 zMobileViTV2Model.forward)T)NNN)rGrHrIrrJr9rrMOBILEVITV2_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr _CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPErrKrLrrrFrMr!r!rBr"rs&   rz MobileViTV2 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for ImageNet. c speZdZeddfdd Zeeeee e e dd e e je ee e je eeee fdddZZS) !MobileViTV2ForImageClassificationNrcs`t||j|_t||_td|jdd}|jdkrJtj||jdnt |_ | dS)Nrrrr)Z in_featuresZ out_features) r8r9 num_labelsrrr#rrrZIdentity classifierr)rAr,r.rBr!r"r9s  z*MobileViTV2ForImageClassification.__init__)rrrr)rrlabelsrrc Csl|dk r |n|jj}|j|||d}|r.|jn|d}||}d}|dk r,|jjdkr|jdkrnd|j_n4|jdkr|jtj ks|jtj krd|j_nd|j_|jjdkrt } |jdkr| | | }n | ||}nN|jjdkrt } | |d|j|d}n|jjdkr,t} | ||}|s\|f|dd} |dk rX|f| S| St|||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). NrrZ regressionZsingle_label_classificationZmulti_label_classificationrmr)losslogitsrk)r,rrrrZ problem_typerZdtyperKlongrrZsqueezerviewrr rk) rArrrroutputsrrrloss_fctrr!r!r"rFs>     "     z)MobileViTV2ForImageClassification.forward)NNNN)rGrHrIrr9rrr_IMAGE_CLASS_CHECKPOINTr r_IMAGE_CLASS_EXPECTED_OUTPUTrrKrLrJrrrFrMr!r!rBr"rs& rcs<eZdZeeeddfdd ZejejdddZZ S)MobileViTV2ASPPPoolingN)r,r-r.rc s4ttjdd|_t|||ddddd|_dS)Nr)rTrsr)r8r9rZAdaptiveAvgPool2d global_poolr+r)rAr,r-r.rBr!r"r9Fs zMobileViTV2ASPPPooling.__init__rDcCs:|jdd}||}||}tjj||ddd}|S)NrbilinearFsizemodeZ align_corners)rrrrrq interpolate)rArEZ spatial_sizer!r!r"rFUs   zMobileViTV2ASPPPooling.forwardrar!r!rBr"rEsrcs<eZdZdZeddfdd ZejejdddZZ S) MobileViTV2ASPPzs ASPP module defined in DeepLab papers: https://arxiv.org/abs/1606.00915, https://arxiv.org/abs/1706.05587 Nrcsttdjdd}|jtjdkr.rd)r8r9r#raspp_out_channelslenZ atrous_ratesr:rrZconvsr+r]extendrprojectrgZaspp_dropout_probr)rAr,Zencoder_out_channelsZ in_projectionZ pool_layerrBrr"r9bs<      zMobileViTV2ASPP.__init__rDcCsDg}|jD]}|||q tj|dd}||}||}|S)Nrrn)rr]rKcatrr)rArEZpyramidconvZpooled_featuresr!r!r"rFs   zMobileViTV2ASPP.forward rGrHrIrWrr9rKrLrFrMr!r!rBr"r]s,rcs<eZdZdZeddfdd ZejejdddZZ S) MobileViTV2DeepLabV3zB DeepLabv3 architecture: https://arxiv.org/abs/1706.05587 Nrc sBtt||_t|j|_t||j |j ddddd|_ dS)NrFT)r-r.r/r4r5r2) r8r9raspprZ Dropout2dZclassifier_dropout_probrr+rrrrAr,rBr!r"r9s  zMobileViTV2DeepLabV3.__init__rjcCs&||d}||}||}|S)Nrm)rrr)rArkrEr!r!r"rFs  zMobileViTV2DeepLabV3.forwardrr!r!rBr"rsrzZ MobileViTV2 model with a semantic segmentation head on top, e.g. for Pascal VOC. c sleZdZeddfdd Zeeeee dd e e j e e j e e e e eeefdddZZS) "MobileViTV2ForSemanticSegmentationNrcs8t||j|_t|dd|_t||_|dS)NF)r)r8r9rrrrsegmentation_headrrrBr!r"r9s   z+MobileViTV2ForSemanticSegmentation.__init__)rr)rrrrrc Cs |dk r |n|jj}|dk r |n|jj}|j|d|d}|rB|jn|d}||}d}|dk r|jjdkrvtdn6tj j ||j ddddd } t |jj d } | | |}|s|r|f|dd} n|f|d d} |dk r|f| S| St|||r|jnddd S) a 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). Returns: Examples: ```python >>> import requests >>> import torch >>> from PIL import Image >>> from transformers import AutoImageProcessor, MobileViTV2ForSemanticSegmentation >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256") >>> model = MobileViTV2ForSemanticSegmentation.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256") >>> inputs = image_processor(images=image, return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**inputs) >>> # logits are of shape (batch_size, num_labels, height, width) >>> logits = outputs.logits ```NTrrz/The number of labels should be greater than onerrFr)Z ignore_indexr)rrrkZ attentions)r,rrrrkrrr:rrqrrrZsemantic_loss_ignore_indexr) rArrrrrZencoder_hidden_statesrrZupsampled_logitsrrr!r!r"rFsB'     z*MobileViTV2ForSemanticSegmentation.forward)NNNN)rGrHrIrr9rrrrrrrKrLrJrrrFrMr!r!rBr"rs   r)rN)=rWtypingrrrrKZtorch.utils.checkpointrZtorch.nnrrrZ activationsr Zmodeling_outputsr r r rZmodeling_utilsrrrrrrrZconfiguration_mobilevitv2rZ get_loggerrGloggerrrrrrZ)MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LISTrr#rr*Moduler+rNrXrbrwrrrrrZMOBILEVITV2_START_DOCSTRINGrrrrrrrr!r!r!r"sd      "A1?)rt  TQ=