o Zhe @sdZddlZddlmZddlmZmZmZmZm Z ddl Z ddl Z ddl m Z ddl mZmZmZddlmZdd lmZmZmZmZmZdd lmZdd lmZmZmZdd lm Z d dl!m"Z"e#e$Z%eGdddeZ&eGdddeZ'eGdddeZ(eGdddeZ)Gddde j*Z+Gddde j*Z,Gddde j*Z-dNde j.d e/d!e0d"e j.fd#d$Z1Gd%d&d&e j*Z2Gd'd(d(e j*Z3Gd)d*d*e j*Z4Gd+d,d,e j*Z5d-e j.d.ee6d/ee6d"e j.fd0d1Z7Gd2d3d3e j*Z8d-e j.d4ee6e6fd5ee6e6fd6eee6d"e j.f d7d8Z9eGd9d:d:eZ:Gd;d<dd>e:ZedCdDGdEdFdFe:Z?edGdDGdHdIdIe:Z@edJdDGdKdLdLe:e ZAgdMZBdS)OzPyTorch Hiera model.N) dataclass)DictListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BackboneOutputBaseModelOutputBaseModelOutputWithPoolingImageClassifierOutput ModelOutput)PreTrainedModel)auto_docstringlogging torch_int) BackboneMixin) HieraConfigc@sreZdZUdZdZeejed<dZ ee ejdfed<dZ ee ejdfed<dZ ee ejdfed<dS)HieraEncoderOutputa Hiera encoder's outputs, with potential hidden states and attentions. 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. 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 stage) of shape `(batch_size, sequence_length, hidden_size)`. Thesre are the unrolled hidden states of the model. 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 stage) 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. reshaped_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 stage) of shape `(batch_size, height, width, hidden_size)`. These are the reshaped and re-rolled hidden states of the model. Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions. Nlast_hidden_state. hidden_states attentionsreshaped_hidden_states) __name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrrrr&r&W/var/www/auris/lib/python3.10/site-packages/transformers/models/hiera/modeling_hiera.pyr+s rc@seZdZUdZdZeejed<dZ eejed<dZ ej ed<dZ eej ed<dZeeejdfed<dZeeejdfed <dZeeejdfed <dS) HieraModelOutputa Hiera model's outputs that also contains a pooling of the last hidden states. 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)`, *optional*, returned when `add_pooling_layer=True` is passed): Average pooling of the last layer hidden-state. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): Tensor indicating which patches are masked (0) and which are not (1). ids_restore (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Tensor containing the original index of the (shuffled) masked patches. 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 stage) of shape `(batch_size, sequence_length, hidden_size)`. These are the unrolled hidden states of the model. 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 stage) 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. reshaped_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 stage) of shape `(batch_size, height, width, hidden_size)`. These are the reshaped and re-rolled hidden states of the model. Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions. Nr pooler_outputbool_masked_pos ids_restore.rrr)rr r!r"rrr#r$r%r)r* BoolTensorr+ LongTensorrrrrr&r&r&r'r(Ls r(c@seZdZUdZdZeejed<dZ eejed<dZ ee ejdfed<dZ ee ejdfed<dZ ee ejdfed<dS) !HieraForImageClassificationOutputa  Hiera image classification outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, `optional`): Loss value for the training task. logits (`torch.FloatTensor` of shape `(batch_size, num_labels)`): Prediction scores of the classification head (logits of the output layer). hidden_states (`tuple(torch.FloatTensor)`, `optional`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. These are the unrolled hidden states of the model. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, `optional`): Tuple of `torch.FloatTensor` (one for each stage) 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. reshaped_hidden_states (`tuple(torch.FloatTensor)`, `optional`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. These are the reshaped and re-rolled hidden states of the model. Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions. Nlosslogits.rrr)rr r!r"r/rr#r$r%r0rrrrr&r&r&r'r.vs r.c@seZdZUdZdZeejed<dZ eejed<dZ ej ed<dZ eej ed<dZeeejed<dZeeejed<dZeeejed <dS) HieraForPreTrainingOutputa Class for HieraForPreTraining's outputs, with potential hidden states and attentions. Args: loss (`torch.FloatTensor` of shape `(1,)`): Pixel reconstruction loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): Tensor indicating which patches are masked (0) and which are not (1). ids_restore (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Tensor containing the original index of the (shuffled) masked patches. 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. reshaped_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, height, width, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to include the spatial dimensions. Nr/r0r*r+rrr)rr r!r"r/rr#r$r%r0r*r,r+r-rrrrr&r&r&r'r1s r1c seZdZdZddeffdd Z ddejdeej d ej fd d Z ddejd eejd e ej ej ffd dZ ddejd eejd e ej eej eej ffddZZS)HieraPatchEmbeddingsz 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. Fis_maecstt|j|_|jdkrtd|jd|j|_|jdd|_ddt|j|j D|_ ddt|j |j D|_ |j |_ ||_tj|j|j|j|j |jd|_dS) NzAThe number of dimensions of the input image should be 2, but got .cSg|]\}}||qSr&r&.0isr&r&r' z1HieraPatchEmbeddings.__init__..cSr7r&r&r8r&r&r'r<r=) kernel_sizestridepadding)super__init__lenZ patch_sizeZ spatial_dims ValueError num_channels image_sizezip patch_stridetokens_spatial_shapemasked_unit_sizemask_spatial_shape mask_ratior3rConv2d embed_dimZ patch_padding projection)selfconfigr3 __class__r&r'rBs"    zHieraPatchEmbeddings.__init__N pixel_valuesr*returncCs^|dur ||S|jdd}|j|jddg|jR}tjj||d}|||S)zZero-out the masked regions of the input before conv. Prevents leakage of masked regions when using overlapping kernels. Nr4rr)size)rOshapeviewrKr functional interpolatefloat)rPrTr*Z target_sizer&r&r' masked_convs  z HieraPatchEmbeddings.masked_convnoisec Cs|jd}t|j}t|d|j}|dur!tj|||jd}tj |dd}tj |dd |j}tj ||g|jd}d|ddd|f<tj |d|d }||fS)a Perform per-sample random masking by per-sample shuffling. Per-sample shuffling is done by argsort random noise. Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`) noise (`torch.FloatTensor` of shape `(batch_size, num_mask_units)`, *optional*) which is mainly used for testing purposes to control randomness and maintain the reproducibility rrNdevicedim)raindex)rWmathprodrKintrLr#randr_Zargsorttozerosgatherbool) rPrTr] batch_size num_windowsZlen_keepZ ids_shuffler+r*r&r&r'random_maskings z#HieraPatchEmbeddings.random_maskingcCsD|jr |j||dnd\}}|||}|ddd}|||fS)Nr])NNr4r)r3rmr\flatten transpose)rPrTr]r*r+ embeddingsr&r&r'forwards   zHieraPatchEmbeddings.forwardFN)rr r!r"rjrBr#r$rr,Tensorr\rr-rmrr __classcell__r&r&rRr'r2s6  %r2c seZdZdZddededdffdd Zd ejd ejd e d e dejf d dZ d ejd e d e dedej f ddZ  ddej de ej dedeeje eje ejffddZZS)HieraEmbeddingsz2 Construct position and patch embeddings. FrQr3rUNcst|j|_ddt|j|jD}ddt||jD|_t||_ ||_ t ||d|_ t td|j |j|_dS)NcSr7r&r&r8r&r&r'r<%r=z,HieraEmbeddings.__init__..cSr7r&r&r8r&r&r'r<&r=r3r)rArBrHrGrFrJrKrcrdZ num_tokensr3r2patch_embeddingsr Parameterr#rhrNposition_embeddings)rPrQr3rIrRr&r'rB"s  zHieraEmbeddings.__init__rq pos_embedsheightwidthc Cs|jd}|jd}tjs||kr||kr|S|jd}||jd}||jd} t|d} |d| | |}|dddd}tj j ||| fddd }|dddd dd|}|S) a2 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, no class embeddings, and different patch strides. 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 rr?r r4ZbicubicF)rVmodeZ align_corners) rWr#Zjit is_tracingrHrreshapepermuterrYrZrX) rPrqr|r}r~Z num_patchesZ num_positionsraZ new_heightZ new_widthZsqrt_num_positionsr&r&r'interpolate_pos_encoding.s$   z(HieraEmbeddings.interpolate_pos_encodingrcCs|r |||j||S|jSrt)rr{)rPrqr}r~rr&r&r'get_position_embeddingTs z&HieraEmbeddings.get_position_embeddingrTr]c CsD|jdd\}}|j||d\}}}||||||}|||fS)Nr6rn)rWryr) rPrTr]rr}r~rqr*r+r&r&r'rr]s zHieraEmbeddings.forwardrsNF)rr r!r"rrjrBr#rurerr$rrrr,r-rrrvr&r&rRr'rwsH  &  rwcs~eZdZdZ   ddedededed ed ed d ffd d Z ddejde ej ded e eje ejffddZ Z S)HieraMaskUnitAttentionz Computes either Mask Unit or Global Attention. Also is able to perform query pooling. Note: this assumes the tokens have already been flattened and unrolled into mask units. rrF hidden_sizehidden_size_output num_heads query_stride window_sizeuse_mask_unit_attnrUNcsbt||_||_||_|||_|jd|_t|d||_ t|||_ ||_ ||_ dS)Ngr ) rArBrrrhead_dimscalerLinearqkvprojrr)rPrrrrrrrRr&r'rBps    zHieraMaskUnitAttention.__init__r head_maskoutput_attentionscCs|j\}}}d}|jr||j|j}||}||d|d|j|j}|dddddd}| d\} } } |jdkrO| ||j||jd|j} | j ddj } | |j | dd } | jdd} |d urh| |} | | } | dd|d|j} || } |r| | fS| d fS) z3Input should be of shape [batch, tokens, channels].rrr rr4r`r6N)rWrrrrrrrrZunbindrXmaxvaluesrrpZsoftmaxrr)rPrrrrkseq_len_rlrquerykeyvalue attn_weightsZ attn_outputr&r&r'rrs&     zHieraMaskUnitAttention.forward)rrFr)rr r!r"rerjrBr#rurr$rrrrvr&r&rRr'ris> rFinput drop_probtrainingrUcCsd|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. rrrr)dtyper_)rWndimr#rfrr_Zfloor_div)rrrZ keep_probrWZ random_tensoroutputr&r&r' drop_paths rcsTeZdZdZd deeddffdd Zdejdejfdd Z de fd d Z Z S) HieraDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).NrrUcst||_dSrt)rArBr)rPrrRr&r'rBs  zHieraDropPath.__init__rcCst||j|jSrt)rrrrPrr&r&r'rrszHieraDropPath.forwardcCs d|jS)Nzp={})formatrrPr&r&r' extra_reprs zHieraDropPath.extra_reprrt) rr r!r"rr[rBr#rurrstrrrvr&r&rRr'rs rcs<eZdZdeddffdd ZdejdejfddZZS) HieraMlprarUNcsJtt|j|_t|t||j|_ tt||j||_ dSrt) rArBr Z hidden_act activation_fnrrreZ mlp_ratiofc1fc2)rPrQrarRr&r'rBs  zHieraMlp.__init__rcCs"||}||}||}|Srt)rrrrr&r&r'rrs   zHieraMlp.forward) rr r!rerBr#rurrrvr&r&rRr'rsrcseZdZ    ddedededed ed ed ed d ffdd Z ddejde ej ded e eje ejffddZ Z S) HieraLayerrrrFrrrrrrrrUNc st||_||_||_tj||jd|_t ||||||d|_ tj||jd|_ t |||_ |dkr9t|nt|_||krKt|||_dSdS)NZeps)rrrrrrr)rArBrrrr LayerNormlayer_norm_epslayernorm_beforerattnlayernorm_afterrmlprIdentityrrr) rPrQrrrrrrrrRr&r'rBs&  zHieraLayer.__init__rrrc Cs|j\}}}||}|j|jkr%||}|||jd|jjddj}|j |||d\}}|| |}|} | |}| |}| | |}||fS)Nrrr`r) rWrrrrrXrrrrrrr) rPrrrrkrrZhidden_states_normrZresidualr&r&r'rrs       zHieraLayer.forward)rrrFr)rr r!rer[rjrBr#rurr$rrrrvr&r&rRr'rsB  %rcseZdZ ddededededeedeeded ed eed dffd d Z dde j dee j ded e e j ee j ffddZ ZS) HieraStageNdepthrrrrrrr stage_numrUc sbtd| durj| dkr| dndtf ddt|D|_dS)NFrrc sDg|]}t|dkr n||po|dkdqS)r)rQrrrrrrr)rr9r: rQrrrrZ$previous_stage_used_masked_attentionrrrr&r'r<6s z'HieraStage.__init__..)rArBmasked_unit_attentionr ModuleListrangelayers) rPrQrrrrrrrrrrRrr'rB s   zHieraStage.__init__FrrrcCsBt|jD]\}}|dur||nd}||||d\}}q||fS)Nr) enumerater)rPrrrr:Z layer_modulelayer_head_maskrr&r&r'rrEs  zHieraStage.forwardrtrs)rr r!rerr[rjrrBr#rur$rrrrvr&r&rRr'rsB     &rrrWmask_unit_shapecCsr|jd|jd}}ddt||D}|j|g|||R}|dddddd }|j|g||R}|S) a] Restore spatial organization by undoing windowed organization of mask units. Args: hidden_states (`torch.Tensor`): The hidden states tensor of shape `[batch_size, num_mask_unit_height*num_mask_unit_width, hidden_size]`. shape (`List[int]`): The original shape of the hidden states tensor before windowing. mask_unit_shape (`List[int]`): The shape of the mask units used for windowing. Returns: torch.Tensor: The restored hidden states tensor of shape [batch_size, num_mask_unit_height*mask_unit_height, num_mask_unit_width*mask_unit_width, hidden_size]. rrcSr7r&r&)r9r;mur&r&r'r<`r=z"undo_windowing..rr r4rr)rWrGrXrr)rrWrrkrnum_mask_unitsr&r&r'undo_windowingQs  rcseZdZdeddffdd Z ddejdedeej dejfd d Z   ddejdeej d eej de de de de eeffddZZS) HieraEncoderrQrUNc sttj}ddtjdj|ddD}tjjddd }|dj fddt |D}t |_j}dg|}tj}tj} tjD]E\} } tjj| } t| || j| ||| || d||| || dt|| | j| | d } | }|j| qWd dtjjD}jgtjdd }i|_t tjD]} ||f|j| <| j krd dt|jD}|dd}qd |_ dS) NcSsg|]}|qSr&)item)r9xr&r&r'r<psz)HieraEncoder.__init__..rcpur^cs$g|]}|vrtjndqSr)rcrdrrrQZquery_pool_layerr&r'r<ts$r) rQrrrrrrrrrcSr7r&r&r8r&r&r'r<r=rcSr7r&r&r8r&r&r'r<r=F)!rArBsumdepthsr#ZlinspaceZdrop_path_rateZtensorZcumsumtolistnum_query_poolrrrstagesrNrcrdrJrrreembed_dim_multiplierrrrappendrGrFrHrCschedulegradient_checkpointing)rPrQZ total_depthZdprZcumulative_depthsZ query_stridesrZ stage_endsZmasked_unit_areaZquery_stride_areaZ idx_stagerrstageZ stage_sizeunroll_schedulerRrr'rBlsJ          zHieraEncoder.__init__r stage_idxr*c Cs|j|\}}|j\}}}t|} dg| } |D]C} |j|g| |t| | |R}|ddddddd}t| D] } | | | | 9<q=|j|dg| |R}|jd}q|j||g| |R}|d urn|St ||| }|S) a\ Roll the given tensor back up to spatial order assuming it's from the given block. If no bool_masked_pos is provided returns: - [batch_size, height, width, hidden_size] If a bool_masked_pos is provided returns: - [batch_size, num_mask_units, mask_unit_height, mask_unit_width, hidden_size] rrr rr4rrN) rrWrCrXrcrdrrrr) rPrrr*rrVrkrrZnum_dimrstridesr:r&r&r'rerolls4      zHieraEncoder.rerollFTrroutput_hidden_states return_dictcCs|rdnd}|r dnd}|rdnd} |r&||f}|j|d|d} || f}t|jD]E\} } |dur7|| nd} |jrI|jrI|| j|| |}n| || |}|d}|r\| |df} |rp||f}|j|| |d} || f}q+|stdd||| |fDSt||| |dS)Nr&r)rr*rcss|] }|dur|VqdSrtr&)r9vr&r&r' sz'HieraEncoder.forward..)rrrr) rrrrrZ_gradient_checkpointing_func__call__tupler)rPrr*rrrrZall_hidden_statesZall_reshaped_hidden_statesZall_self_attentionsrr:Z stage_modulerZ layer_outputsr&r&r'rrs@          zHieraEncoder.forwardrt)NNFFT)rr r!rrBr#rurerr,rr$rjrrrrrrvr&r&rRr'rks@3 1 r image_shaperHrc Cs|j\}}}ddt||D}|}|j|g||g}|D]W} ddt|| D}ddt|| D} |g| |g} || }t| } dgttd| ddttd| dd| dg} || }|dt| }|t | 9}q| dt ||}|S) a Reorders the tokens such that patches are contiguous in memory. E.g., given [batch_size, (height, width), hidden_size] and stride of (stride, stride), this will re-order the tokens as [batch_size, (stride, stride, height // stride, width // stride), hidden_size] This allows operations like Max2d to be computed as x.view(batch_size, stride*stride, -1, hidden_size).max(dim=1). Not only is this faster, but it also makes it easy to support inputs of arbitrary dimensions in addition to patch-wise sparsity. Performing this operation multiple times in sequence puts entire windows as contiguous in memory. For instance, if you applied the stride (2, 2) 3 times, entire windows of size 8x8 would be contiguous in memory, allowing operations like mask unit attention computed easily and efficiently, while also allowing max to be applied sequentially. Note: This means that intermediate values of the model are not in height x width order, so they need to be re-rolled if you want to use the intermediate values as a height x width feature map. The last block of the network is fine though, since by then the strides are all consumed. cSr7r&r&r8r&r&r'r<r=zunroll..cSr7r&r&r8r&r&r'r<!r=cSsg|] }|D]}|qqSr&r&)r9pairrr&r&r'r<#srr4rr) rWrGrXrClistrrrorcrdr) rrrHrrkrrrVZ current_sizerZ new_shapeZnum_dimsrr&r&r'unrolls  8 rc@s&eZdZeZdZdZdZdddZdS) HieraPreTrainedModelhierarTTrUNcCs|jj}t|trtjj|j|ddSt|tr-tjj|j |dtjj|j |ddSt|tj tj tj frRtjj|j|d|jdurPtj|j|dSdSt|tjrltj|j|tj|j|jjdSdS)zInitialize the weights)stdN)rQZinitializer_range isinstancerwrinitZ trunc_normal_r{ HieraDecoder mask_tokendecoder_position_embeddingsrZConv1drMweightZbiasZ constant_rZlayer_norm_init)rPmodulerr&r&r' _init_weights=s    z"HieraPreTrainedModel._init_weights)rUN) rr r!rZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrr&r&r&r'r6s rcs8eZdZdeffdd ZdejdejfddZZS) HieraPoolerrQcsJtt|j|jt|jd}tj||j d|_ t d|_ dS)Nrr) rArBrerNrrCrrrr layernormZAdaptiveAvgPool1dpooler)rPrQ num_featuresrRr&r'rBSs zHieraPooler.__init__rrUcCs0|dd}||}t|d}||}|S)Nrr4)rprr#ror)rPr pooled_outputr&r&r'rrYs    zHieraPooler.forward) rr r!rrBr#rurrrvr&r&rRr'rRsrcseZdZddededeffdd Zdefd d Zd ee e e fdd fd dZ e dde ejde ejde ejde ede ede ede edeeeffddZZS) HieraModelTFrQadd_pooling_layerr3cs~t|t|j|jt|jd|_t||d|_ t ||_ |j gt|jdd|_ |r6t|nd|_|dS)z add_pooling_layer (`bool`, *optional*, defaults to `True`): Whether or not to apply pooling layer. is_mae (`bool`, *optional*, defaults to `False`): Whether or not to run the model on MAE mode. rrxNr)rArBrerNrrCrrrwrqrencoderrrrr post_init)rPrQrr3rRr&r'rBcs    zHieraModel.__init__rUcC|jjSrtrqryrr&r&r'get_input_embeddingswzHieraModel.get_input_embeddingsheads_to_pruneNcCs*|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)itemsrlayerZ attentionZ prune_heads)rPrrZheadsr&r&r' _prune_headszszHieraModel._prune_headsrTr]rrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td||t|jj}|j|||d\}} } |j d|j df} t || |jj |j d} | durvt |jj} | j \}}}| dd| |}| |} | |d|} |j| | ||||d}|d }d}|jdur||}|s|dur||fn|f}| dur|| | fn|}||ddSt||| | |j|j|jd S) z noise (`torch.FloatTensor` of shape `(batch_size, num_mask_units)`, *optional*): Mainly used for testing purposes to control randomness and maintain the reproducibility Nz You have to specify pixel_values)rr]r6r)rrHrr)r*rrrrr)rr)r*r+rrr)rQrruse_return_dictrDZ get_head_maskrCrrqrWrrHrrcrdrJZ unsqueezeZtilerXrrr(rrr)rPrTr]rrrrrembedding_outputr*r+rrZmask_unit_arearkrrZ positionsZencoder_outputsZsequence_outputrZ head_outputsr&r&r'rrsb    zHieraModel.forward)TFNNNNNNN)rr r!rrjrBr2rrrerrrrr#rur$rrrrrrvr&r&rRr'ras: rc s\eZdZdeffdd Z  d dejdejdeejd e d e ejejff d d Z Z S)rrQc sPttjjtjd}ddtjj D}fddt|j D|_ fddtj j D|_ t|j|_ttddj|_ttdt|j j|_tjjjjddgjdgjdd |_tjjjd |_j d j d j |_!|j!tj j"}tj||_#dS) NrcSr7r&r&r8r&r&r'r<r=z)HieraDecoder.__init__..cg|] \}}||jqSr&rr8rQr&r'r<cr r&r r8r r&r'r<r Frr) rQrrrrrrrrrr)$rArBrerNrrCrrGrFrHrtokens_spatial_shape_finalrJmask_unit_spatial_shape_finalrrdecoder_hidden_sizedecoder_embeddingsrzr#rhrrcrdrrZdecoder_num_headsZ decoder_depth decoder_blockrr decoder_normr pred_striderE decoder_pred)rPrQrrIZpred_dimrRr r'rBs:        zHieraDecoder.__init__NFencoder_hidden_statesr*rrrUc Cs6||}|jdd\}}}|j\} } tj| | ||||j|jd} |jddddd} || | ddd}| dd|||}| | |<d| | | | } t | |j |j}t |dddf|j |j}||jdd|jd}||jdd}||j}|j|||d\}} ||}||}||fS)Nr4)r_rrr.r)rr)rrWr#rhr_rrrXrexpandror[rrrrrrr)rPrr*rrrmask_unit_heightmask_unit_widthrrkrZdecoder_hidden_statesZ mask_tokensrr&r&r'rrsP         zHieraDecoder.forwardr) rr r!rrBr#rur,rrjrrrrvr&r&rRr'rs+rcsXeZdZdeffdd ZdejdejdejfddZ d e ejdejfd d Z Z S) HieraMultiScaleHeadrQc stfddtjjD|_fddttjD|_ j}t |_ tj D]+}ddt||jD}ddt|jD}|j t j|j ||j d||dq0|j t dS)Ncr r&r r8r r&r'r<@r z0HieraMultiScaleHead.__init__..c g|] }tjj|qSr&rerNrrr r&r'r<CcSr7r&r&r8r&r&r'r<Jr=cSr7r&r&r8r&r&r'r<Kr=r)r>r?)rArBrGrJrrrrCrZstage_dimensionsrrmulti_scale_fusion_headsrrrMr)rPrQZcurrent_masked_unit_sizeidxZkernelrRr r'rB>s*      zHieraMultiScaleHead.__init__headrrUc Cst|tjr|S|j\}}}}}||||||}|dddd}||}|dddd}|jdd\}} }||||| |}|S)Nrr rr4)rrrrWrr) rPr rrkrrrrZmask_unit_height_finalZmask_unit_width_finalr&r&r'apply_fusion_headVs   z%HieraMultiScaleHead.apply_fusion_head feature_mapscCs.d}t|j|D] \}}||||}q|S)Nr)rGrr!)rPr"rr Z feature_mapr&r&r'rrmszHieraMultiScaleHead.forward) rr r!rrBrModuler#rur!rrrrvr&r&rRr'r=s"ra; The Hiera Model transformer with the decoder on top for self-supervised pre-training. 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 Zdejdejdejfdd Zdejd ejdejfd d Z e       dde ejd e ej de ejde e de e de e de e deeeffddZZS)HieraForPreTrainingrQrUNcs\t|t|ddd|_tj|jj|jd|_t ||_ t ||_ |j j |_ |dS)NFTrr3r)rArBrrrrrr encoder_normrmultiscale_fusionrdecoderrrrPrQrRr&r'rBs     zHieraForPreTraining.__init__rTr*cCs|dddd}|j}|d||d||}|ddd}||}|jjr@|jddd}|jddd}|||dd }|S) Nrr4r rrT)raZkeepdimgư>r)rrZunfoldrorQZnormalize_pixel_lossmeanvar)rPrTr*rVlabelr*r+r&r&r'get_pixel_label_2dsz&HieraForPreTraining.get_pixel_label_2dr0cCs2|}|||}||}||d}|}|S)Nr4)r-r*)rPrTr0r*r,r/r&r&r' forward_losss   z HieraForPreTraining.forward_lossr]rrrrrc CsN|dur|n|jj}|dur|n|jj}|dur|n|jj}|j||||d||d}|d} |d} |d} | d|jjjd| df} || } || } |j| | ||d\} } | || | }|s| | | f}|rr||df}|r{||d f}|r||df}|dur|f|S|St || | | |r|j nd|j |r|j d Sdd S) a noise (`torch.FloatTensor` of shape `(batch_size, num_mask_units)`, *optional*): Mainly used for testing purposes to control randomness and maintain the reproducibility Examples: ```python >>> from transformers import AutoImageProcessor, HieraForPreTraining >>> 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/hiera-tiny-224-mae-hf") >>> model = HieraForPreTraining.from_pretrained("facebook/hiera-tiny-224-mae-hf") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> loss = outputs.loss >>> print(list(logits.shape)) [1, 196, 768] ```NT)r]rrrrrrrr4)r*rrr r)r/r0r*r+rrr)rQrrrrrr'r&r(r.r1rrr)rPrTr]rrrrroutputsr"r*Zids_to_restoreZfused_hidden_statesr0r/rr&r&r'rrs\$       zHieraForPreTraining.forwardr )rr r!rrBr#rur,r-r.rrr$rjrrr1rrrvr&r&rRr'r$vs:   r$a Hiera Model transformer with an image classification head on top (a linear layer on top of the final hidden state with average pooling) e.g. for ImageNet. Note that it's possible to fine-tune Hiera on higher resolution images than the ones it has been trained on, by setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained position embeddings to the higher resolution. cseZdZdeddffdd Ze      ddeejdeejdee d ee d ee d ee de e e ffd d Z ZS)HieraForImageClassificationrQrUNcsVt||j|_t|ddd|_|jdkr t|jj|jnt|_ | dS)NTFr%r) rArB num_labelsrrrrrr classifierrr)rRr&r'rBs " z$HieraForImageClassification.__init__rlabelsrrrrcCs|dur|n|jj}|dur|n|jj}|dur|n|jj}|j||||||d}|d} || } d} |dur|| j}|jjdurh|j dkrNd|j_n|j dkrd|j t j ks_|j t j krdd|j_nd|j_|jjdkrt} |j dkr| | |} n+| | |} n%|jjdkrt} | | d|j |d} n|jjdkrt} | | |} |s| f|dd} | dur| f| S| St| | |j|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). N)rrrrrrZ regressionZsingle_label_classificationZmulti_label_classificationrr4)r/r0rrr)rQrrrrr2rgr_Z problem_typer1rr#longrer Zsqueezer rXr r.rrr)rPrTrr3rrrrr/rr0r/Zloss_fctrr&r&r'rr#sX     "       z#HieraForImageClassification.forward)NNNNNN)rr r!rrBrrr#rurjrrr.rrrvr&r&rRr'r0s0 r0zN Hiera backbone, to be used with frameworks like DETR and MaskFormer. c s^eZdZdeffdd ZddZ   ddejdee d ee d ee d e f d d Z Z S) HieraBackbonerQcsttjgfddttjD|_tdd|_ t |_ i}t |j |jD] \}}t|||<q4t||_|dS)Ncrr&rrr r&r'r<wrz*HieraBackbone.__init__..Frx)rArBZ_init_backbonerNrrCrrrwrqrrrGZ _out_featuresZchannelsrrZ ModuleDicthidden_states_normsr)rPrQr6rrErRr r'rBss       zHieraBackbone.__init__cCrrtrrr&r&r'rrz"HieraBackbone.get_input_embeddingsNrTrrrrUcCsD|dur|n|jj}|dur|n|jj}|dur|n|jj}||\}}}|j|d|d|d}|d}d} t|j|D]7\} } | |jvrs| j \} } }}| | | ||} |j | | } | | | ||} | dddd  } | | f7} q<|s| f}|r||df7}|r||d f7}|St| |r|dnd|r|d d Sdd S) a? Returns: 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("facebook/hiera-tiny-224-hf") >>> model = AutoBackbone.from_pretrained( ... "facebook/hiera-tiny-224-hf", 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, 7, 7] ```NT)rrrrrr&rr rr4)r"rr)rQrrrrqrrGZ stage_namesZ out_featuresrWrXr6r contiguousr)rPrTrrrrrr/rr"rZ hidden_staterkr}r~rErr&r&r'rrsJ    zHieraBackbone.forward)NNN) rr r!rrBrr#rurrjrrrrvr&r&rRr'r5ms"r5)r0r$r5rr)rF)Cr"rc dataclassesrtypingrrrrrr#Ztorch.utils.checkpointrZtorch.nnr r r Z activationsr Zmodeling_outputsrrrrrZmodeling_utilsrutilsrrrZutils.backbone_utilsrZconfiguration_hierarZ get_loggerrloggerrr(r.r1r#r2rwrrur[rjrrrrrrerrrrrrrrr$r0r5__all__r&r&r&r's        )#$^L E@&2    6ri9 Y b