





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_lowerCamelCase      							=   logging.get_logger(__name__)

# General docstring
_lowerCamelCase      							=   'ResNetConfig'

# Base docstring
_lowerCamelCase      							=   'microsoft/resnet-50'
_lowerCamelCase      							=   [1, 2048, 7, 7]

# Image classification docstring
_lowerCamelCase      							=   'microsoft/resnet-50'
_lowerCamelCase      							=   'tiger cat'

_lowerCamelCase      							=   [
    'microsoft/resnet-50',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__ = 3   ,				a__ = 1   ,				a__ = "relu"):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  Optional[int]						=		nn.Convad(
                 a__   ,				a__   ,				kernel_size=a__   ,				stride=a__   ,				padding=kernel_size // 2   ,				bias=a__)
             _lowerCamelCase				:  Optional[int]						=		nn.BatchNormad(a__)
             _lowerCamelCase				:  List[Any]						=		ACTaFN[activation] if activation is not None else nn.Identity()



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  int						=		self.convolution(a__)
             _lowerCamelCase				:  int						=		self.normalization(a__)
             _lowerCamelCase				:  Optional[int]						=		self.activation(a__)
             return hidden_state







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  List[str]						=		ResNetConvLayer(
                 config.num_channels   ,				config.embedding_size   ,				kernel_size=7   ,				stride=2   ,				activation=config.hidden_act)
             _lowerCamelCase				:  Union[str, Any]						=		nn.MaxPoolad(kernel_size=3   ,				stride=2   ,				padding=1)
             _lowerCamelCase				:  Optional[int]						=		config.num_channels



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  int						=		pixel_values.shape[1]
             if num_channels != self.num_channels:
                    raise ValueError(
                        '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.''')
             _lowerCamelCase				:  Any						=		self.embedder(a__)
             _lowerCamelCase				:  Optional[Any]						=		self.pooler(a__)
             return embedding







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__ = 2):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  Union[str, Any]						=		nn.Convad(a__   ,				a__   ,				kernel_size=1   ,				stride=a__   ,				bias=a__)
             _lowerCamelCase				:  str						=		nn.BatchNormad(a__)



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		self.convolution(a__)
             _lowerCamelCase				:  Any						=		self.normalization(a__)
             return hidden_state







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__ = 1   ,				a__ = "relu"):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  List[str]						=		in_channels != out_channels or stride != 1
             _lowerCamelCase				:  List[str]						=		(
                 ResNetShortCut(a__   ,				a__   ,				stride=a__) if should_apply_shortcut else nn.Identity()
             )
             _lowerCamelCase				:  Tuple						=		nn.Sequential(
                 ResNetConvLayer(a__   ,				a__   ,				stride=a__)   ,				ResNetConvLayer(a__   ,				a__   ,				activation=a__)   ,				)
             _lowerCamelCase				:  Tuple						=		ACTaFN[activation]



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  int						=		hidden_state
             _lowerCamelCase				:  Optional[Any]						=		self.layer(a__)
             _lowerCamelCase				:  Optional[Any]						=		self.shortcut(a__)
             hidden_state += residual
             _lowerCamelCase				:  List[str]						=		self.activation(a__)
             return hidden_state







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__ = 1   ,				a__ = "relu"   ,				a__ = 4):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  Optional[Any]						=		in_channels != out_channels or stride != 1
             _lowerCamelCase				:  int						=		out_channels // reduction
             _lowerCamelCase				:  Optional[Any]						=		(
                 ResNetShortCut(a__   ,				a__   ,				stride=a__) if should_apply_shortcut else nn.Identity()
             )
             _lowerCamelCase				:  Union[str, Any]						=		nn.Sequential(
                 ResNetConvLayer(a__   ,				a__   ,				kernel_size=1)   ,				ResNetConvLayer(a__   ,				a__   ,				stride=a__)   ,				ResNetConvLayer(a__   ,				a__   ,				kernel_size=1   ,				activation=a__)   ,				)
             _lowerCamelCase				:  Union[str, Any]						=		ACTaFN[activation]



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		hidden_state
             _lowerCamelCase				:  Dict						=		self.layer(a__)
             _lowerCamelCase				:  Optional[int]						=		self.shortcut(a__)
             hidden_state += residual
             _lowerCamelCase				:  Union[str, Any]						=		self.activation(a__)
             return hidden_state







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__   ,				a__   ,				a__   ,				a__ = 2   ,				a__ = 2   ,				):



             """simple docstring"""





             super().__init__()

             _lowerCamelCase				:  str						=		ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

             _lowerCamelCase				:  List[Any]						=		nn.Sequential(
                 # downsampling is done in the first layer with stride of 2
                 layer(a__   ,				a__   ,				stride=a__   ,				activation=config.hidden_act)   ,				*[layer(a__   ,				a__   ,				activation=config.hidden_act) for _ in range(depth - 1)]   ,				)



      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  str						=		input
             for layer in self.layers:
                    _lowerCamelCase				:  List[Any]						=		layer(a__)
             return hidden_state







class       __A							(   nn.Module	):




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             super().__init__()
             _lowerCamelCase				:  Union[str, Any]						=		nn.ModuleList([])
             # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
             self.stages.append(
                 ResNetStage(
                     a__   ,				config.embedding_size   ,				config.hidden_sizes[0]   ,				stride=2 if config.downsample_in_first_stage else 1   ,				depth=config.depths[0]   ,				))
             _lowerCamelCase				:  List[Any]						=		zip(config.hidden_sizes   ,				config.hidden_sizes[1:])
             for (in_channels, out_channels), depth in zip(a__   ,				config.depths[1:]):
                    self.stages.append(ResNetStage(a__   ,				a__   ,				a__   ,				depth=a__))



      def 					__snake_case      (		self   ,				a__   ,				a__ = False   ,				a__ = True):



             """simple docstring"""





             _lowerCamelCase				:  Dict						=		() if output_hidden_states else None

             for stage_module in self.stages:
                    if output_hidden_states:
                           _lowerCamelCase				:  Optional[Any]						=		hidden_states + (hidden_state,)

                    _lowerCamelCase				:  Union[str, Any]						=		stage_module(a__)

             if output_hidden_states:
                    _lowerCamelCase				:  Dict						=		hidden_states + (hidden_state,)

             if not return_dict:
                    return tuple(v for v in [hidden_state, hidden_states] if v is not None)

             return BaseModelOutputWithNoAttention(
                 last_hidden_state=a__   ,				hidden_states=a__   ,				)







class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      UpperCAmelCase__							=   ResNetConfig
      UpperCAmelCase__							=   """resnet"""
      UpperCAmelCase__							=   """pixel_values"""
      UpperCAmelCase__							=   True


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             if isinstance(a__   ,				nn.Convad):
                    nn.init.kaiming_normal_(module.weight   ,				mode='''fan_out'''   ,				nonlinearity='''relu''')
             elif isinstance(a__   ,				(nn.BatchNormad, nn.GroupNorm)):
                    nn.init.constant_(module.weight   ,				1)
                    nn.init.constant_(module.bias   ,				0)



      def 					__snake_case      (		self   ,				a__   ,				a__=False):



             """simple docstring"""





             if isinstance(a__   ,				a__):
                    _lowerCamelCase				:  Tuple						=		value
_lowerCamelCase      							=   r'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

_lowerCamelCase      							=   r'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'







@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""			,lowerCamelCase__			,)
class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             super().__init__(a__)
             _lowerCamelCase				:  Tuple						=		config
             _lowerCamelCase				:  Optional[int]						=		ResNetEmbeddings(a__)
             _lowerCamelCase				:  Optional[int]						=		ResNetEncoder(a__)
             _lowerCamelCase				:  int						=		nn.AdaptiveAvgPoolad((1, 1))
             # Initialize weights and apply final processing
             self.post_init()



      @add_start_docstrings_to_model_forward(a__)
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC   ,				output_type=a__   ,				config_class=_CONFIG_FOR_DOC   ,				modality='''vision'''   ,				expected_output=_EXPECTED_OUTPUT_SHAPE   ,				)
      def 					__snake_case      (		self   ,				a__   ,				a__ = None   ,				a__ = None):



             """simple docstring"""





             _lowerCamelCase				:  Tuple						=		(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _lowerCamelCase				:  List[Any]						=		return_dict if return_dict is not None else self.config.use_return_dict

             _lowerCamelCase				:  Union[str, Any]						=		self.embedder(a__)

             _lowerCamelCase				:  Union[str, Any]						=		self.encoder(
                 a__   ,				output_hidden_states=a__   ,				return_dict=a__)

             _lowerCamelCase				:  List[Any]						=		encoder_outputs[0]

             _lowerCamelCase				:  Optional[int]						=		self.pooler(a__)

             if not return_dict:
                    return (last_hidden_state, pooled_output) + encoder_outputs[1:]

             return BaseModelOutputWithPoolingAndNoAttention(
                 last_hidden_state=a__   ,				pooler_output=a__   ,				hidden_states=encoder_outputs.hidden_states   ,				)







@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """			,lowerCamelCase__			,)
class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             super().__init__(a__)
             _lowerCamelCase				:  Any						=		config.num_labels
             _lowerCamelCase				:  int						=		ResNetModel(a__)
             # classification head
             _lowerCamelCase				:  str						=		nn.Sequential(
                 nn.Flatten()   ,				nn.Linear(config.hidden_sizes[-1]   ,				config.num_labels) if config.num_labels > 0 else nn.Identity()   ,				)
             # initialize weights and apply final processing
             self.post_init()



      @add_start_docstrings_to_model_forward(a__)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT   ,				output_type=a__   ,				config_class=_CONFIG_FOR_DOC   ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,				)
      def 					__snake_case      (		self   ,				a__ = None   ,				a__ = None   ,				a__ = None   ,				a__ = None   ,				):



             """simple docstring"""





             _lowerCamelCase				:  Dict						=		return_dict if return_dict is not None else self.config.use_return_dict

             _lowerCamelCase				:  Optional[int]						=		self.resnet(a__   ,				output_hidden_states=a__   ,				return_dict=a__)

             _lowerCamelCase				:  str						=		outputs.pooler_output if return_dict else outputs[1]

             _lowerCamelCase				:  Tuple						=		self.classifier(a__)

             _lowerCamelCase				:  Union[str, Any]						=		None

             if labels is not None:
                    if self.config.problem_type is None:
                           if self.num_labels == 1:
                                  _lowerCamelCase				:  Dict						=		'''regression'''
                           elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                                  _lowerCamelCase				:  Tuple						=		'''single_label_classification'''
                           else:
                                  _lowerCamelCase				:  Any						=		'''multi_label_classification'''
                    if self.config.problem_type == "regression":
                           _lowerCamelCase				:  int						=		MSELoss()
                           if self.num_labels == 1:
                                  _lowerCamelCase				:  Dict						=		loss_fct(logits.squeeze()   ,				labels.squeeze())
                           else:
                                  _lowerCamelCase				:  List[Any]						=		loss_fct(a__   ,				a__)
                    elif self.config.problem_type == "single_label_classification":
                           _lowerCamelCase				:  int						=		CrossEntropyLoss()
                           _lowerCamelCase				:  Optional[int]						=		loss_fct(logits.view(-1   ,				self.num_labels)   ,				labels.view(-1))
                    elif self.config.problem_type == "multi_label_classification":
                           _lowerCamelCase				:  int						=		BCEWithLogitsLoss()
                           _lowerCamelCase				:  List[str]						=		loss_fct(a__   ,				a__)

             if not return_dict:
                    _lowerCamelCase				:  List[Any]						=		(logits,) + outputs[2:]
                    return (loss,) + output if loss is not None else output

             return ImageClassifierOutputWithNoAttention(loss=a__   ,				logits=a__   ,				hidden_states=outputs.hidden_states)







@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """			,lowerCamelCase__			,)
class       __A							(   lowerCamelCase__			,lowerCamelCase__	):




      """simple docstring"""







      def __init__(		self   ,				a__):



             """simple docstring"""





             super().__init__(a__)
             super()._init_backbone(a__)

             _lowerCamelCase				:  str						=		[config.embedding_size] + config.hidden_sizes
             _lowerCamelCase				:  Optional[int]						=		ResNetEmbeddings(a__)
             _lowerCamelCase				:  Optional[Any]						=		ResNetEncoder(a__)

             # initialize weights and apply final processing
             self.post_init()



      @add_start_docstrings_to_model_forward(a__)
      @replace_return_docstrings(output_type=a__   ,				config_class=_CONFIG_FOR_DOC)
      def 					__snake_case      (		self   ,				a__   ,				a__ = None   ,				a__ = None):



             """simple docstring"""





             _lowerCamelCase				:  Any						=		return_dict if return_dict is not None else self.config.use_return_dict
             _lowerCamelCase				:  Optional[int]						=		(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )

             _lowerCamelCase				:  Dict						=		self.embedder(a__)

             _lowerCamelCase				:  Any						=		self.encoder(a__   ,				output_hidden_states=a__   ,				return_dict=a__)

             _lowerCamelCase				:  Optional[int]						=		outputs.hidden_states

             _lowerCamelCase				:  int						=		()
             for idx, stage in enumerate(self.stage_names):
                    if stage in self.out_features:
                           feature_maps += (hidden_states[idx],)

             if not return_dict:
                    _lowerCamelCase				:  int						=		(feature_maps,)
                    if output_hidden_states:
                           output += (outputs.hidden_states,)
                    return output

             return BackboneOutput(
                 feature_maps=a__   ,				hidden_states=outputs.hidden_states if output_hidden_states else None   ,				attentions=a__   ,				)








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase      							=   logging.get_logger(__name__)

_lowerCamelCase      							=   {
    'sayakpaul/vit-msn-base': 'https://huggingface.co/sayakpaul/vit-msn-base/resolve/main/config.json',
    # See all ViT MSN models at https://huggingface.co/models?filter=vit_msn
}







class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      UpperCAmelCase__							=   """vit_msn"""


      def __init__(		self   ,				a__=768   ,				a__=12   ,				a__=12   ,				a__=3072   ,				a__="gelu"   ,				a__=0.0   ,				a__=0.0   ,				a__=0.02   ,				a__=1e-06   ,				a__=224   ,				a__=16   ,				a__=3   ,				a__=True   ,				**a__   ,				):



             """simple docstring"""





             super().__init__(**a__)

             _lowerCamelCase				:  Tuple						=		hidden_size
             _lowerCamelCase				:  str						=		num_hidden_layers
             _lowerCamelCase				:  int						=		num_attention_heads
             _lowerCamelCase				:  Optional[int]						=		intermediate_size
             _lowerCamelCase				:  Union[str, Any]						=		hidden_act
             _lowerCamelCase				:  str						=		hidden_dropout_prob
             _lowerCamelCase				:  List[Any]						=		attention_probs_dropout_prob
             _lowerCamelCase				:  Optional[int]						=		initializer_range
             _lowerCamelCase				:  Optional[Any]						=		layer_norm_eps
             _lowerCamelCase				:  Dict						=		image_size
             _lowerCamelCase				:  Tuple						=		patch_size
             _lowerCamelCase				:  List[str]						=		num_channels
             _lowerCamelCase				:  Optional[int]						=		qkv_bias

