a h @sdZddlZddlmZddlmZddlmZm Z ddl Z ddl Z ddl Z ddl mZddlmZdd lmZdd lmZdd lmZmZdd lmZmZdd lmZddlmZmZddl m!Z!m"Z"m#Z#m$Z$ddl%m&Z&m'Z'ddl(m)Z)m*Z*ddl+m,Z,e$-e.Z/ee#ddGddde!Z0ee#ddGddde!Z1ddZ2Gdddej3Z4Gdd d ej3Z5d@ej3e j6e j6e j6e e j6e7e7d"d#d$Z8Gd%d&d&ej3Z9Gd'd(d(ej3Z:Gd)d*d*ej3Z;Gd+d,d,ej3ZGd1d2d2ej3Z?e#Gd3d4d4eZ@e#Gd5d6d6e@ZAGd7d8d8ej3ZBe#d9dGd:d;d;e@ZCe#dd>e@ZDgd?ZEdS)Az,PyTorch VideoMAE (masked autoencoder) model.N)deepcopy) dataclass)CallableOptional)nn)MSELoss)ACT2FN)GradientCheckpointingLayer)BaseModelOutputImageClassifierOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputTransformersKwargsauto_docstringlogging)IMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STD)can_return_tuplecheck_model_inputs)VideoMAEConfigz[ Class for VideoMAEDecoder's outputs, with potential hidden states and attentions. )Z custom_introc@sPeZdZUdZdZeejed<dZ ee ejed<dZ ee ejed<dS)VideoMAEDecoderOutputz logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. Nlogits hidden_states attentions) __name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rtuplerr(r(j/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/videomae/modeling_videomae.pyr+s rzb Class for VideoMAEForPreTraining's outputs, with potential hidden states and attentions. c@sbeZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)VideoMAEForPreTrainingOutputz loss (`torch.FloatTensor` of shape `(1,)`): Pixel reconstruction loss. logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. Nlossrrr) r r!r"r#r+rr$r%r&rrr'rr(r(r(r)r*<s r*csfddtfddt|D}t|dddddf|dddddf<t|dddddf|dddddf<t|dS) z Sinusoid position encoding tablecsfddtDS)Nc s(g|] }tdd|dqS)i')nppower).0Zhid_j)d_hidpositionr(r) WzOget_sinusoid_encoding_table..get_position_angle_vec..)ranger1)r0r5r)get_position_angle_vecVsz;get_sinusoid_encoding_table..get_position_angle_veccsg|] }|qSr(r()r/Zpos_i)r6r(r)r2Yr3z/get_sinusoid_encoding_table..Nrr,r)r-arrayr4sincosr$r%Z unsqueeze)Z n_positionr0Zsinusoid_tabler()r0r6r)get_sinusoid_encoding_tableRs  ..r:cs(eZdZdZfddZddZZS)VideoMAEEmbeddingsz7 Construct the patch and position embeddings. cs8tt||_|jj|_t|j|j|_||_dSN) super__init__VideoMAEPatchEmbeddingspatch_embeddings num_patchesr: hidden_sizeposition_embeddingsconfigselfrD __class__r(r)r>fs    zVideoMAEEmbeddings.__init__cCsZ||}||j|j|jdd}|durV|j\}}}||}||d|}|S)NTdevicecopy)r@rCdetachtype_astorJshapereshape)rF pixel_valuesbool_masked_pos embeddings batch_size_ num_channelsr(r(r)forwardos   zVideoMAEEmbeddings.forwardr r!r"r#r>rX __classcell__r(r(rGr)r;`s r;cs(eZdZdZfddZddZZS)r?aw Video to Patch Embedding. This module turns a batch of videos of shape (batch_size, num_frames, num_channels, height, width) into a tensor of shape (batch_size, seq_len, hidden_size) to be consumed by a Transformer encoder. The seq_len (the number of patches) equals (number of frames // tubelet_size) * (height // patch_size) * (width // patch_size). c st|j}|j}|j}|j}|j}|j}t|t j j r@|n||f}t|t j j rZ|n||f}||_||_t ||_|d|d|d|d||j}||_||_ tj|||j|d|df|j|d|dfd|_dS)Nrr)Z in_channelsZ out_channelsZ kernel_sizeZstride)r=r> image_size patch_sizerWrB num_frames tubelet_size isinstance collectionsabcIterableintrArConv3d projection) rFrDr[r\rWrBr]r^rArGr(r)r>s,  (z VideoMAEPatchEmbeddings.__init__c Cs|j\}}}}}||jkr"td||jdks>||jdkrltd|d|d|jdd|jdd |dddd d }||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,r)rPrW ValueErrorr[permutereflatten transpose)rFrRrUr]rWheightwidthrTr(r(r)rXs (zVideoMAEPatchEmbeddings.forwardrYr(r(rGr)r?s r?)modulequerykeyvalueattention_maskscalingdropoutc Ks|t||dd|}tjj|dtjd|j}tjj |||j d}|durX||}t||} | dd } | |fS)NrL)dimdtype)ptrainingrr,) r$matmulrkr functionalZsoftmaxZfloat32rOrxrurz contiguous) rorprqrrrsrtrukwargsZ attn_weightsZ attn_outputr(r(r)eager_attention_forwards  rcsHeZdZeddfdd ZdeejeejejfdddZ Z S) VideoMAESelfAttentionN)rDreturncst|j|jdkr>t|ds>td|jd|jd||_|j|_t|j|j|_|j|j|_ |j |_ |jd|_ d|_ tj|j|j dd|_tj|j|j dd|_tj|j|j dd|_|jrtt|j |_tt|j |_n d|_d|_dS) NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .gFbias)r=r>rBnum_attention_headshasattrrhrDrcattention_head_size all_head_sizeZattention_probs_dropout_prob dropout_probrt is_causalrLinearrprqrrZqkv_bias Parameterr$zerosq_biasv_biasrErGr(r)r>s,   zVideoMAESelfAttention.__init__) head_maskrc Cs6|j\}}}|jdur&tj|jddnd}tjj||jj |d}tjj||j j |jd}tjj||j j |jd} | |d|j |jdd} | |d|j |jdd} | |d|j |jdd} t} |jjdkrt|jj} | || | | ||j|j|jsdn|jd \}}|dd |jf}||}||fS) NF)Z requires_grad)inputweightrrLrr,eagerrn)rrtrurv)rPrr$Z zeros_likerrr|ZlinearrqrrrrpviewrrrkrrDZ_attn_implementationr rrtrzrsizerrQ)rFrrrUZ seq_lengthrVZk_biaskeysvaluesZqueriesZ key_layerZ value_layerZ query_layerZattention_interfaceZ context_layerZattention_probsZnew_context_layer_shaper(r(r)rXs0     zVideoMAESelfAttention.forward)N) r r!r"rr>rr$Tensorr'rXrZr(r(rGr)rsrcs>eZdZdZedfdd ZejejejdddZZ S)VideoMAESelfOutputz The residual connection is defined in VideoMAELayer instead of here (as is the case with other models), due to the layernorm applied before each block. rDcs.tt|j|j|_t|j|_dSr<) r=r>rrrBdenseDropouthidden_dropout_probrurErGr(r)r>s zVideoMAESelfOutput.__init__r input_tensorrcCs||}||}|Sr<rrurFrrr(r(r)rXs  zVideoMAESelfOutput.forward) r r!r"r#rr>r$rrXrZr(r(rGr)rsrcsReZdZedfdd ZeedddZd ej e ej ej dd d Z Z S) VideoMAEAttentionrcs*tt||_t||_t|_dSr<)r=r>r attentionroutputset pruned_headsrErGr(r)r>%s   zVideoMAEAttention.__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)Nrrrw)lenrrrrrrrprqrrrrrunion)rFrindexr(r(r) prune_heads+s zVideoMAEAttention.prune_headsNrrrcCs |||\}}|||}|Sr<)rr)rFrrZself_attn_outputrVrr(r(r)rX=s zVideoMAEAttention.forward)N) r r!r"rr>rrcrr$rrrXrZr(r(rGr)r$srcs6eZdZedfdd ZejejdddZZS)VideoMAEIntermediatercsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr<) r=r>rrrBintermediate_sizerr_Z hidden_actstrr intermediate_act_fnrErGr(r)r>Es   zVideoMAEIntermediate.__init__)rrcCs||}||}|Sr<)rr)rFrr(r(r)rXMs  zVideoMAEIntermediate.forward r r!r"rr>r$rrXrZr(r(rGr)rDsrcs:eZdZedfdd ZejejejdddZZS)VideoMAEOutputrcs.tt|j|j|_t|j|_ dSr<) r=r>rrrrBrrrrurErGr(r)r>Us zVideoMAEOutput.__init__rcCs ||}||}||}|Sr<rrr(r(r)rXZs  zVideoMAEOutput.forwardrr(r(rGr)rTsrcsDeZdZdZedfdd Zd ejeejejdddZ Z S) VideoMAELayerz?This corresponds to the Block class in the timm implementation.rcsbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dS)Nreps)r=r>Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr LayerNormrBlayer_norm_epslayernorm_beforelayernorm_afterrErGr(r)r>es    zVideoMAELayer.__init__NrcCsB||}|||}||}||}||}|||}|Sr<)rrrrr)rFrrZhidden_states_normZattention_outputZ layer_outputr(r(r)rXos     zVideoMAELayer.forward)N) r r!r"r#rr>r$rrrXrZr(r(rGr)rbs rcs>eZdZedfdd ZdejeejedddZ Z S) VideoMAEEncoderrcs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr(rr/rVrr(r)r2r3z,VideoMAEEncoder.__init__..F) r=r>rDr ModuleListr4num_hidden_layerslayergradient_checkpointingrErGrr)r>s  zVideoMAEEncoder.__init__NrcCs<t|jD]&\}}|dur"||nd}|||}q t|dS)Nlast_hidden_state) enumeraterr )rFrri layer_moduleZlayer_head_maskr(r(r)rXs zVideoMAEEncoder.forward)N) r r!r"rr>r$rrr rXrZr(r(rGr)rsrc@sDeZdZUeed<dZdZdZdZdZ dZ dZ e e dZddZdS) VideoMAEPreTrainedModelrDvideomaerRT)rrcCsjt|tjtjfr@|jjjd|jjd|j durf|j j n&t|tj rf|j j |jj ddS)zInitialize the weightsrn)meanstdNg?) r_rrrdrdataZnormal_rDZinitializer_rangerZzero_rZfill_)rFror(r(r) _init_weightss   z%VideoMAEPreTrainedModel._init_weightsN)r r!r"rr&Zbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_supports_sdpaZ_supports_flash_attnZ_supports_flex_attnZ_supports_attention_backendrrZ_can_record_outputsrr(r(r(r)rs rc s^eZdZfddZddZddZeed ej e ej e ej e eedd d ZZS) VideoMAEModelcsTt|||_t||_t||_|jr4d|_nt j |j |j d|_| dS)Nr)r=r>rDr;rTrencoderuse_mean_pooling layernormrrrBr post_initrErGr(r)r>s   zVideoMAEModel.__init__cCs|jjSr<)rTr@)rFr(r(r)get_input_embeddingssz"VideoMAEModel.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)rFZheads_to_prunerrr(r(r) _prune_headsszVideoMAEModel._prune_headsNrRrSrr~rcKsN|||jj}|||}|j||d}|j}|jdurD||}t|dS)a bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Each video in the batch must have the same number of masked patches. If `None`, then all patches are considered. Sequence length is `(num_frames // tubelet_size) * (image_size // patch_size) ** 2`. Examples: ```python >>> import av >>> import numpy as np >>> from transformers import AutoImageProcessor, VideoMAEModel >>> from huggingface_hub import hf_hub_download >>> np.random.seed(0) >>> def read_video_pyav(container, indices): ... ''' ... Decode the video with PyAV decoder. ... Args: ... container (`av.container.input.InputContainer`): PyAV container. ... indices (`list[int]`): List of frame indices to decode. ... Returns: ... result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3). ... ''' ... frames = [] ... container.seek(0) ... start_index = indices[0] ... end_index = indices[-1] ... for i, frame in enumerate(container.decode(video=0)): ... if i > end_index: ... break ... if i >= start_index and i in indices: ... frames.append(frame) ... return np.stack([x.to_ndarray(format="rgb24") for x in frames]) >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len): ... ''' ... Sample a given number of frame indices from the video. ... Args: ... clip_len (`int`): Total number of frames to sample. ... frame_sample_rate (`int`): Sample every n-th frame. ... seg_len (`int`): Maximum allowed index of sample's last frame. ... Returns: ... indices (`list[int]`): List of sampled frame indices ... ''' ... converted_len = int(clip_len * frame_sample_rate) ... end_idx = np.random.randint(converted_len, seg_len) ... start_idx = end_idx - converted_len ... indices = np.linspace(start_idx, end_idx, num=clip_len) ... indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64) ... return indices >>> # video clip consists of 300 frames (10 seconds at 30 FPS) >>> file_path = hf_hub_download( ... repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset" ... ) >>> container = av.open(file_path) >>> # sample 16 frames >>> indices = sample_frame_indices(clip_len=16, frame_sample_rate=1, seg_len=container.streams.video[0].frames) >>> video = read_video_pyav(container, indices) >>> image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base") >>> model = VideoMAEModel.from_pretrained("MCG-NJU/videomae-base") >>> # prepare video for the model >>> inputs = image_processor(list(video), return_tensors="pt") >>> # forward pass >>> outputs = model(**inputs) >>> last_hidden_states = outputs.last_hidden_state >>> list(last_hidden_states.shape) [1, 1568, 768] ```rNr)Z get_head_maskrDrrTrrrr )rFrRrSrr~Zembedding_outputZencoder_outputssequence_outputr(r(r)rXs^   zVideoMAEModel.forward)NN)r r!r"r>rrrrr$r%r BoolTensorrrrr rXrZr(r(rGr)rs rcs4eZdZedfdd ZejedddZZ S)VideoMAEDecoderrcst|j|j|jd}t||j_|j_ |j _ |j _ tfddt|jD|_t|j|_|dkrt|j|nt|_d|_|_dS)Nr,csg|] }tqSr(rrZdecoder_configr(r)r2>r3z,VideoMAEDecoder.__init__..rF)r=r>rWr^r\rdecoder_hidden_sizerBZdecoder_num_hidden_layersrZdecoder_num_attention_headsrZdecoder_intermediate_sizerrrr4decoder_layersrnormrIdentityheadrrD)rFrDZdecoder_num_labelsrGrr)r>3s zVideoMAEDecoder.__init__)rreturn_token_numcCsT|jD]}||dd}q|dkr6|dd| df}||}||}t|dS)Nrr)r)rrrr)rFrrrrr(r(r)rXIs   zVideoMAEDecoder.forward) r r!r"rr>r$rrcrXrZr(r(rGr)r2srzb The VideoMAE Model transformer with the decoder on top for self-supervised pre-training. c sJeZdZfddZeedejeje ej e e e dddZZS)VideoMAEForPreTrainingcstt|||_t||_tj|j|jdd|_ t t dd|j|_ t|jjj|j|_t||_|dS)NFrr)r=r>rDrrrrrBrencoder_to_decoderrr$r mask_tokenr:rTrArCrdecoderrrErGr(r)r>^s    zVideoMAEForPreTraining.__init__Nrc Ks&|j|f||d|}|j}||}|j\}}} |durDtd|j|dd|} | j |j dd} | | |d| } | | |d| } t j || |j| gdd} || | jd}|j}d}t |jjd kr|}nd|j }|j}t tj ||d ddddddf}t tj ||d ddddddf}|||}|j\}}} }}|jj|jj}}|jjr@|||||| ||||||}|d dd d dddd }|||||||||||| }||jddd|jdddd d}|||||||||||| }n|jjd krVtd|||||| ||||||}|d dd d dddd }|||||||||||| }|j\}}} || |d| }Wdn1s0Yt!}|||}t"|||j#|j$dS)a bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Each video in the batch must have the same number of masked patches. Sequence length is `(num_frames // tubelet_size) * (image_size // patch_size) ** 2`. Examples: ```python >>> from transformers import AutoImageProcessor, VideoMAEForPreTraining >>> import numpy as np >>> import torch >>> num_frames = 16 >>> video = list(np.random.randint(0, 256, (num_frames, 3, 224, 224))) >>> image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base") >>> model = VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base") >>> pixel_values = image_processor(video, return_tensors="pt").pixel_values >>> num_patches_per_frame = (model.config.image_size // model.config.patch_size) ** 2 >>> seq_length = (num_frames // model.config.tubelet_size) * num_patches_per_frame >>> bool_masked_pos = torch.randint(0, 2, (1, seq_length)).bool() >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) >>> loss = outputs.loss ```)rSrNz!One must provided a boolean mask rLTrIrrr)rJrxrrgr,rv)rwkeepdim)rwZunbiasedrgư>zQCan't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False.r+rrr)%rrrrPrhrCexpandrNrMrOrJrQr$catrrrZno_gradrDrWrxZ as_tensorrrr^r\Z norm_pix_lossrrir}rvarsqrtrr*rr)rFrRrSrr~outputsrrUrVrWZexpanded_position_embeddingsZpos_emb_visibleZ pos_emb_maskZx_fullZdecoder_outputsrr+framesrJrxrrtimerlrmr^r\Z frames_normZ videos_patchlabelsZloss_fctr(r(r)rXos$    &&      2 zVideoMAEForPreTraining.forward)N)r r!r"r>rrr$r%rrrrrr*rXrZr(r(rGr)rXs rz VideoMAE Model transformer with a video classification head on top (a linear layer on top of the average pooled hidden states of all tokens) e.g. for ImageNet. c sReZdZfddZeedeejeejeeje e e dddZ Z S)VideoMAEForVideoClassificationcsft||j|_t||_|jr0t|jnd|_ |jdkrPt |j|jnt |_ | dS)Nr)r=r>Z num_labelsrrrrrrBfc_normrr classifierrrErGr(r)r> s   $z'VideoMAEForVideoClassification.__init__N)rRrrr~rc Ks|j|fd|i|}|j}|jdur<|d}||}n|dddf}||}d} |durz|j|||jfi|} t| ||j|j dS)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). Examples: ```python >>> import av >>> import torch >>> import numpy as np >>> from transformers import AutoImageProcessor, VideoMAEForVideoClassification >>> from huggingface_hub import hf_hub_download >>> np.random.seed(0) >>> def read_video_pyav(container, indices): ... ''' ... Decode the video with PyAV decoder. ... Args: ... container (`av.container.input.InputContainer`): PyAV container. ... indices (`list[int]`): List of frame indices to decode. ... Returns: ... result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3). ... ''' ... frames = [] ... container.seek(0) ... start_index = indices[0] ... end_index = indices[-1] ... for i, frame in enumerate(container.decode(video=0)): ... if i > end_index: ... break ... if i >= start_index and i in indices: ... frames.append(frame) ... return np.stack([x.to_ndarray(format="rgb24") for x in frames]) >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len): ... ''' ... Sample a given number of frame indices from the video. ... Args: ... clip_len (`int`): Total number of frames to sample. ... frame_sample_rate (`int`): Sample every n-th frame. ... seg_len (`int`): Maximum allowed index of sample's last frame. ... Returns: ... indices (`list[int]`): List of sampled frame indices ... ''' ... converted_len = int(clip_len * frame_sample_rate) ... end_idx = np.random.randint(converted_len, seg_len) ... start_idx = end_idx - converted_len ... indices = np.linspace(start_idx, end_idx, num=clip_len) ... indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64) ... return indices >>> # video clip consists of 300 frames (10 seconds at 30 FPS) >>> file_path = hf_hub_download( ... repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset" ... ) >>> container = av.open(file_path) >>> # sample 16 frames >>> indices = sample_frame_indices(clip_len=16, frame_sample_rate=1, seg_len=container.streams.video[0].frames) >>> video = read_video_pyav(container, indices) >>> image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics") >>> model = VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics") >>> inputs = image_processor(list(video), return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**inputs) ... logits = outputs.logits >>> # model predicts one of the 400 Kinetics-400 classes >>> predicted_label = logits.argmax(-1).item() >>> print(model.config.id2label[predicted_label]) eating spaghetti ```rNrrr) rrrrrZ loss_functionrDr rr) rFrRrrr~rrrrr+r(r(r)rXs \    z&VideoMAEForVideoClassification.forward)NNN)r r!r"r>rrrr$rrrr rXrZr(r(rGr)rs r)rrrr)rn)Fr#collections.abcr`rKr dataclassesrtypingrrnumpyr-r$Ztorch.utils.checkpointrZtorch.nnrZ activationsr Zmodeling_layersr Zmodeling_outputsr r Zmodeling_utilsr rZprocessing_utilsrZ pytorch_utilsrrutilsrrrrZutils.constantsrrZ utils.genericrrZconfiguration_videomaerZ get_loggerr loggerrr*r:Moduler;r?rfloatrrrrrrrrrrrrr__all__r(r(r(r)s          != = &'