o i\@sdZddlmZmZddlZddlmZddlm Z m Z m Z ddl m Z mZmZmZmZddlmZdd lmZmZmZmZd d lmZeeZd Zd ZgdZ d Z!dZ"Gdddej#j$Z%Gdddej#j$Z&Gdddej#j$Z'Gdddej#j$Z(Gdddej#j$Z)Gdddej#j$Z*Gdddej#j$Z+Gddde Z,d Z-d!Z.eGd"d#d#ej#j$Z/ed$e-Gd%d&d&e,Z0ed'e-Gd(d)d)e,eZ1gd*Z2dS)+zTensorFlow ResNet model.)OptionalUnionN)ACT2FN) TFBaseModelOutputWithNoAttention*TFBaseModelOutputWithPoolingAndNoAttention&TFImageClassifierOutputWithNoAttention)TFPreTrainedModelTFSequenceClassificationLosskeraskeras_serializable unpack_inputs) shape_list)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardlogging) ResNetConfigrzmicrosoft/resnet-50)rirz tiger catc szeZdZ   ddededededed d f fd d Zd ejd ejfddZdd ejde d ejfddZ dddZ Z S)TFResNetConvLayerrrrelu in_channels out_channels kernel_sizestride activationreturnNc sxtjd i||d|_tjj|||dddd|_tjjdddd |_|dur-t |ntj d |_ ||_ ||_ dS) NvalidF convolution)rstridespaddinguse_biasnameh㈵>? normalizationepsilonmomentumr$linear)super__init__ pad_valuer layersConv2DconvBatchNormalizationr'r Activationrrr)selfrrrrrkwargs __class__r,a/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/resnet/modeling_tf_resnet.pyr.6s   zTFResNetConvLayer.__init__ hidden_statecCs2|j|jf}}t|d||dg}||}|S)N)rr)r/tfpadr2)r5r: height_pad width_padr,r,r9r Js zTFResNetConvLayer.convolutionFtrainingcCs&||}|j||d}||}|SNr?)r r'r)r5r:r?r,r,r9callQs  zTFResNetConvLayer.callcC|jrdSd|_t|dddur2t|jj|jddd|jgWdn1s-wYt|dddur_t|jj|jddd|j gWddS1sXwYdSdS)NTr2r') builtgetattrr; name_scoper2r$buildrr'rr5 input_shaper,r,r9rGW"zTFResNetConvLayer.build)rrrFN) __name__ __module__ __qualname__intstrr.r;Tensorr boolrBrG __classcell__r,r,r7r9r5s(rcsPeZdZdZdeddffdd Zddejd edejfd d Z dd d Z Z S)TFResNetEmbeddingszO ResNet Embeddings (stem) composed of a single aggressive convolution. configrNc sPtjd i|t|j|jdd|jdd|_tjj ddddd|_ |j|_dS) Nrrembedder)rrrr$rrpooler) pool_sizer!r"r$r,) r-r.r num_channelsembedding_size hidden_actrWr r0 MaxPool2DrXr5rVr6r7r,r9r.hs zTFResNetEmbeddings.__init__F pixel_valuesr?cCsjt|\}}}}tr||jkrtd|}||}t|ddgddgddgddgg}||}|S)NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rr)rr;executing_eagerlyrZ ValueErrorrWr<rX)r5r_r?_rZr:r,r,r9rBus $ zTFResNetEmbeddings.callcC|jrdSd|_t|dddur-t|jj|jdWdn1s(wYt|dddurUt|jj|jdWddS1sNwYdSdS)NTrWrX)rDrEr;rFrWr$rGrXrHr,r,r9rG"zTFResNetEmbeddings.buildrKrL) rMrNrO__doc__rr.r;rRrSrBrGrTr,r,r7r9rUcs   rUc sZeZdZdZddedededdffdd Zdd ejd edejfd dZ dddZ Z S)TFResNetShortCutz ResNet shortcut, used to project the residual features to the correct size. If needed, it is also used to downsample the input using `stride=2`. rrrrrNc sNtjd i|tjj|d|ddd|_tjjdddd|_||_||_ dS) NrFr )rr!r#r$r%r&r'r(r,) r-r.r r0r1r r3r'rr)r5rrrr6r7r,r9r.s  zTFResNetShortCut.__init__Fxr?cCs |}||}|j||d}|Sr@)r r')r5rgr?r:r,r,r9rBs zTFResNetShortCut.callcCrC)NTr r') rDrEr;rFr r$rGrr'rrHr,r,r9rGrJzTFResNetShortCut.build)rrKrL) rMrNrOrerPr.r;rRrSrBrGrTr,r,r7r9rfs   rfc s`eZdZdZ ddededededd f fd d Zdd ejde dejfddZ dddZ Z S)TFResNetBasicLayerzO A classic ResNet's residual layer composed by two `3x3` convolutions. rrrrrrrNc sztjd i|||kp|dk}t|||dd|_t||ddd|_|r-t|||ddntjjddd|_ t ||_ dS) Nrlayer.0rr$layer.1rr$shortcutr+r$r,) r-r.rconv1conv2rfr r0r4rmrr)r5rrrrr6should_apply_shortcutr7r,r9r.szTFResNetBasicLayer.__init__Fr:r?cCsD|}|j||d}|j||d}|j||d}||7}||}|Sr@)rorprmrr5r:r?residualr,r,r9rBs zTFResNetBasicLayer.callcCs|jrdSd|_t|dddur-t|jj|jdWdn1s(wYt|dddurRt|jj|jdWdn1sMwYt|dddurzt|jj|jdWddS1sswYdSdS)NTrorprm) rDrEr;rFror$rGrprmrHr,r,r9rGs "zTFResNetBasicLayer.build)rrrKrL rMrNrOrerPrQr.r;rRrSrBrGrTr,r,r7r9rhs  rhc sheZdZdZ   ddedededed ed d f fd d Zddejde d ejfddZ dddZ Z S)TFResNetBottleNeckLayera% A classic ResNet's bottleneck layer composed by three `3x3` convolutions. The first `1x1` convolution reduces the input by a factor of `reduction` in order to make the second `3x3` convolution faster. The last `1x1` convolution remaps the reduced features to `out_channels`. rrrrrr reductionrNc  stjd i|||kp|dk}||}t||ddd|_t|||dd|_t||dddd|_|r;t|||ddntjj d dd |_ t ||_ dS) Nrri)rr$rkrjzlayer.2)rrr$rmr+rnr,) r-r.rconv0rorprfr r0r4rmrr) r5rrrrrwr6rqreduces_channelsr7r,r9r.s z TFResNetBottleNeckLayer.__init__Fr:r?cCsR|}|j||d}|j||d}|j||d}|j||d}||7}||}|Sr@)rxrorprmrrrr,r,r9rBs zTFResNetBottleNeckLayer.callcCsB|jrdSd|_t|dddur-t|jj|jdWdn1s(wYt|dddurRt|jj|jdWdn1sMwYt|dddurwt|jj|jdWdn1srwYt|dddurt|j j|j dWddS1swYdSdS)NTrxrorprm) rDrEr;rFrxr$rGrorprmrHr,r,r9rGs("zTFResNetBottleNeckLayer.build)rrrvrKrLrtr,r,r7r9rus(  ruc sdeZdZdZ ddedededededd f fd d Zdd ejde dejfddZ dddZ Z S) TFResNetStagez4 A ResNet stage composed of stacked layers. rrVrrrdepthrNc sftjdi|jdkrtnt||jddg}|fddt|dD7}||_dS)N bottleneckzlayers.0)rrr$c s(g|]}jd|ddqS)zlayers.rrl)r\).0irVlayerrr,r9 !sz*TFResNetStage.__init__..rr,)r-r. layer_typerurhr\range stage_layers)r5rVrrrr{r6r0r7rr9r.s  zTFResNetStage.__init__Fr:r?cCs|jD]}|||d}q|Sr@)r)r5r:r?rr,r,r9rB's zTFResNetStage.callc Cj|jrdSd|_t|dddur1|jD]}t|j |dWdn1s+wYqdSdS)NTr)rDrErr;rFr$rGr5rIrr,r,r9rG,  zTFResNetStage.build)rrrKrL) rMrNrOrerrPr.r;rRrSrBrGrTr,r,r7r9rzs$rzc sXeZdZdeddffdd Z   ddejd ed ed edef d d Z dddZ Z S)TFResNetEncoderrVrNc stjdi|t||j|jd|jrdnd|jdddg|_tt |j|jdd|jddD]\}\}}}|j t||||d|ddq4dS) Nrrrzstages.0)rr{r$zstages.)r{r$r,) r-r.rzr[ hidden_sizesdownsample_in_first_stagedepthsstages enumeratezipappend)r5rVr6r~rrr{r7r,r9r.7s  &zTFResNetEncoder.__init__FTr:output_hidden_states return_dictr?cCsf|rdnd}|jD]}|r||f}|||d}q |r ||f}|s-tdd||fDSt||dS)Nr,rAcss|] }|dur|VqdSrLr,)r}vr,r,r9 \z'TFResNetEncoder.call..)last_hidden_state hidden_states)rtupler)r5r:rrr?r stage_moduler,r,r9rBIs     zTFResNetEncoder.callc Cr)NTr)rDrErr;rFr$rGrr,r,r9rG`rzTFResNetEncoder.build)FTFrL) rMrNrOrr.r;rRrSrrBrGrTr,r,r7r9r6s" rc@s(eZdZdZeZdZdZeddZ dS)TFResNetPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. resnetr_cCs dtjd|jjddftjdiS)Nr_)shapedtype)r; TensorSpecrVrZfloat32)r5r,r,r9input_signaturets z'TFResNetPreTrainedModel.input_signatureN) rMrNrOrer config_classbase_model_prefixmain_input_namepropertyrr,r,r,r9rjsrad This model is a TensorFlow [keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and behavior. Parameters: config ([`ResNetConfig`]): 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 [`~TFPreTrainedModel.from_pretrained`] method to load the model weights. a> Args: pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`ConvNextImageProcessor.__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. csveZdZeZdeddffdd Ze   ddejde e d e e d e de e eje ff d d Zdd dZZS)TFResNetMainLayerrVrNc sHtjdi|||_t|dd|_t|dd|_tjj dd|_ dS)NrWrnencoderT)keepdimsr,) r-r.rVrUrWrrr r0GlobalAveragePooling2DrXr^r7r,r9r.s zTFResNetMainLayer.__init__Fr_rrr?c Cs|dur|n|jj}|dur|n|jj}tj|gdd}|j||d}|j||||d}|d}||}t|d}t|d}d} |ddD] } | td d | D} qJ|s`||f| S|rd| nd} t ||| d S) N)rrrr)permrArrr?rrrrrr,rcss|] }t|dVqdS)rN)r; transpose)r}hr,r,r9rrz)TFResNetMainLayer.call..)r pooler_outputr) rVruse_return_dictr;rrWrrXrr) r5r_rrr?embedding_outputencoder_outputsr pooled_outputrr:r,r,r9rBs.      zTFResNetMainLayer.callcCrc)NTrWr)rDrEr;rFrWr$rGrrHr,r,r9rGrdzTFResNetMainLayer.buildNNFrL)rMrNrOrrr.r r;rRrrSrrrrBrGrTr,r,r7r9rs& -rzOThe bare ResNet model outputting raw features without any specific head on top.cseZdZdeddffdd Zeeeee e de de   dd e jd eed eed edeee je ff d dZdddZZS) TFResNetModelrVrNc s&tj|fi|t|dd|_dS)Nr)rVr$)r-r.rrr^r7r,r9r.szTFResNetModel.__init__vision) checkpoint output_typermodalityexpected_outputFr_rrr?cCs>|dur|n|jj}|dur|n|jj}|j||||d}|S)N)r_rrr?)rVrrr)r5r_rrr?resnet_outputsr,r,r9rBszTFResNetModel.callcCsd|jrdSd|_t|dddur0t|jj|jdWddS1s)wYdSdS)NTr)rDrEr;rFrr$rGrHr,r,r9rGs"zTFResNetModel.buildrrL)rMrNrOrr.rRESNET_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPEr r;rRrrSrrrBrGrTr,r,r7r9rs4rz ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for ImageNet. cseZdZdeddffdd ZdejdejfddZee e e e e ed e     dd eejd eejd eedeededeeeje ff ddZdddZZS)TFResNetForImageClassificationrVrNc sbtj|fi||j|_t|dd|_|jdkr#tjj|jddntjjddd|_ ||_ dS)Nrrnrz classifier.1r+) r-r. num_labelsrrr r0Denser4classifier_layerrVr^r7r,r9r.s  z'TFResNetForImageClassification.__init__rgcCstj|}||}|SrL)r r0Flattenr)r5rglogitsr,r,r9 classifiers z)TFResNetForImageClassification.classifier)rrrrFr_labelsrrr?c Cs|dur|n|jj}|j||||d}|r|jn|d}||}|dur'dn|||} |sC|f|dd} | durA| f| S| St| ||jdS)a) labels (`tf.Tensor` 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 classification loss is computed (Cross-Entropy). Nrrr)lossrr)rVrrrrhf_compute_lossrr) r5r_rrrr?outputsrrroutputr,r,r9rB!s z#TFResNetForImageClassification.callcCs|jrdSd|_t|dddur-t|jj|jdWdn1s(wYt|dddur\t|jj|jdd|jj dgWddS1sUwYdSdS)NTrr) rDrEr;rFrr$rGrrVrrHr,r,r9rGHs"z$TFResNetForImageClassification.build)NNNNFrL)rMrNrOrr.r;rRrrrr_IMAGE_CLASS_CHECKPOINTrr_IMAGE_CLASS_EXPECTED_OUTPUTr rrSrrrBrGrTr,r,r7r9rs< r)rrr)3retypingrr tensorflowr;activations_tfrmodeling_tf_outputsrrrmodeling_tf_utilsr r r r r tf_utilsrutilsrrrrconfiguration_resnetr get_loggerrMloggerrrrrrr0LayerrrUrfrhrurzrrRESNET_START_DOCSTRINGrrrr__all__r,r,r,r9sL    .*"+:"4 D+ E