
'''simple docstring'''



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


_lowercase    =		logging.get_logger(__name__)

# General docstring
_lowercase    =		"""ResNetConfig"""

# Base docstring
_lowercase    =		"""microsoft/resnet-50"""
_lowercase    =		[1, 2048, 7, 7]

# Image classification docstring
_lowercase    =		"""microsoft/resnet-50"""
_lowercase    =		"""tiger cat"""

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




class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






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

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			nn.Convad(
												    _lowercase   ,       _lowercase   ,       kernel_size=_lowercase   ,       stride=_lowercase   ,       padding=kernel_size // 2   ,       bias=_lowercase				)
												_lowerCAmelCase  						=			nn.BatchNormad(_lowercase				)
												_lowerCAmelCase  						=			ACTaFN[activation] if activation is not None else nn.Identity()




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.convolution(_lowercase				)
												_lowerCAmelCase  						=			self.normalization(_lowercase				)
												_lowerCAmelCase  						=			self.activation(_lowercase				)
												return hidden_state






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase				):

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			ResNetConvLayer(
												    config.num_channels   ,       config.embedding_size   ,       kernel_size=7   ,       stride=2   ,       activation=config.hidden_act				)
												_lowerCAmelCase  						=			nn.MaxPoolad(kernel_size=3   ,       stride=2   ,       padding=1				)
												_lowerCAmelCase  						=			config.num_channels




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			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  						=			self.embedder(_lowercase				)
												_lowerCAmelCase  						=			self.pooler(_lowercase				)
												return embedding






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






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

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			nn.Convad(_lowercase   ,       _lowercase   ,       kernel_size=1   ,       stride=_lowercase   ,       bias=_lowercase				)
												_lowerCAmelCase  						=			nn.BatchNormad(_lowercase				)




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.convolution(_lowercase				)
												_lowerCAmelCase  						=			self.normalization(_lowercase				)
												return hidden_state






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






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

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			in_channels != out_channels or stride != 1
												_lowerCAmelCase  						=			(
												    ResNetShortCut(_lowercase   ,       _lowercase   ,       stride=_lowercase				) if should_apply_shortcut else nn.Identity()
												)
												_lowerCAmelCase  						=			nn.Sequential(
												    ResNetConvLayer(_lowercase   ,       _lowercase   ,       stride=_lowercase				)   ,       ResNetConvLayer(_lowercase   ,       _lowercase   ,       activation=_lowercase				)   ,       )
												_lowerCAmelCase  						=			ACTaFN[activation]




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			hidden_state
												_lowerCAmelCase  						=			self.layer(_lowercase				)
												_lowerCAmelCase  						=			self.shortcut(_lowercase				)
												hidden_state += residual
												_lowerCAmelCase  						=			self.activation(_lowercase				)
												return hidden_state






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






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

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			in_channels != out_channels or stride != 1
												_lowerCAmelCase  						=			out_channels // reduction
												_lowerCAmelCase  						=			(
												    ResNetShortCut(_lowercase   ,       _lowercase   ,       stride=_lowercase				) if should_apply_shortcut else nn.Identity()
												)
												_lowerCAmelCase  						=			nn.Sequential(
												    ResNetConvLayer(_lowercase   ,       _lowercase   ,       kernel_size=1				)   ,       ResNetConvLayer(_lowercase   ,       _lowercase   ,       stride=_lowercase				)   ,       ResNetConvLayer(_lowercase   ,       _lowercase   ,       kernel_size=1   ,       activation=_lowercase				)   ,       )
												_lowerCAmelCase  						=			ACTaFN[activation]




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			hidden_state
												_lowerCAmelCase  						=			self.layer(_lowercase				)
												_lowerCAmelCase  						=			self.shortcut(_lowercase				)
												hidden_state += residual
												_lowerCAmelCase  						=			self.activation(_lowercase				)
												return hidden_state






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






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

												"""simple docstring"""



												super().__init__()

												_lowerCAmelCase  						=			ResNetBottleNeckLayer if config.layer_type == """bottleneck""" else ResNetBasicLayer

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




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			input
												for layer in self.layers:
																		_lowerCAmelCase  						=			layer(_lowercase				)
												return hidden_state






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase				):

												"""simple docstring"""



												super().__init__()
												_lowerCAmelCase  						=			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(
												        _lowercase   ,       config.embedding_size   ,       config.hidden_sizes[0]   ,       stride=2 if config.downsample_in_first_stage else 1   ,       depth=config.depths[0]   ,       )				)
												_lowerCAmelCase  						=			zip(config.hidden_sizes   ,       config.hidden_sizes[1:]				)
												for (in_channels, out_channels), depth in zip(_lowercase   ,       config.depths[1:]				):
																		self.stages.append(ResNetStage(_lowercase   ,       _lowercase   ,       _lowercase   ,       depth=_lowercase				)				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase = False   ,       _lowercase = True				):

												"""simple docstring"""



												_lowerCAmelCase  						=			() if output_hidden_states else None

												for stage_module in self.stages:
																		if output_hidden_states:
																								_lowerCAmelCase  						=			hidden_states + (hidden_state,)

																		_lowerCAmelCase  						=			stage_module(_lowercase				)

												if output_hidden_states:
																		_lowerCAmelCase  						=			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=_lowercase   ,       hidden_states=_lowercase   ,       )






class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						_lowercase		:				Optional[int]									=       ResNetConfig
						_lowercase		:				str									=       '''resnet'''
						_lowercase		:				List[str]									=       '''pixel_values'''
						_lowercase		:				Optional[int]									=       True




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



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




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase=False				):

												"""simple docstring"""



												if isinstance(_lowercase   ,       _lowercase				):
																		_lowerCAmelCase  						=			value






_lowercase    =		R"""
    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 ([`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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_lowercase    =		R"""
    Args:
        pixel_values (`torch.FloatTensor` 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.
"""




@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''	,							_SCREAMING_SNAKE_CASE	,							)
class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase				):

												"""simple docstring"""



												super().__init__(_lowercase				)
												_lowerCAmelCase  						=			config
												_lowerCAmelCase  						=			ResNetEmbeddings(_lowercase				)
												_lowerCAmelCase  						=			ResNetEncoder(_lowercase				)
												_lowerCAmelCase  						=			nn.AdaptiveAvgPoolad((1, 1)				)
												# Initialize weights and apply final processing
												self.post_init()




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

												"""simple docstring"""



												_lowerCAmelCase  						=			(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												_lowerCAmelCase  						=			return_dict if return_dict is not None else self.config.use_return_dict

												_lowerCAmelCase  						=			self.embedder(_lowercase				)

												_lowerCAmelCase  						=			self.encoder(
												    _lowercase   ,       output_hidden_states=_lowercase   ,       return_dict=_lowercase				)

												_lowerCAmelCase  						=			encoder_outputs[0]

												_lowerCAmelCase  						=			self.pooler(_lowercase				)

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

												return BaseModelOutputWithPoolingAndNoAttention(
												    last_hidden_state=_lowercase   ,       pooler_output=_lowercase   ,       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.
    '''	,							_SCREAMING_SNAKE_CASE	,							)
class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase				):

												"""simple docstring"""



												super().__init__(_lowercase				)
												_lowerCAmelCase  						=			config.num_labels
												_lowerCAmelCase  						=			ResNetModel(_lowercase				)
												# classification head
												_lowerCAmelCase  						=			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(_lowercase				)
						@add_code_sample_docstrings(
						    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,       output_type=_lowercase   ,       config_class=_CONFIG_FOR_DOC   ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,       )
						def  _lowercase     ( self   ,       _lowercase = None   ,       _lowercase = None   ,       _lowercase = None   ,       _lowercase = None   ,       ):

												"""simple docstring"""



												_lowerCAmelCase  						=			return_dict if return_dict is not None else self.config.use_return_dict

												_lowerCAmelCase  						=			self.resnet(_lowercase   ,       output_hidden_states=_lowercase   ,       return_dict=_lowercase				)

												_lowerCAmelCase  						=			outputs.pooler_output if return_dict else outputs[1]

												_lowerCAmelCase  						=			self.classifier(_lowercase				)

												_lowerCAmelCase  						=			None

												if labels is not None:
																		if self.config.problem_type is None:
																								if self.num_labels == 1:
																														_lowerCAmelCase  						=			"""regression"""
																								elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																														_lowerCAmelCase  						=			"""single_label_classification"""
																								else:
																														_lowerCAmelCase  						=			"""multi_label_classification"""
																		if self.config.problem_type == "regression":
																								_lowerCAmelCase  						=			MSELoss()
																								if self.num_labels == 1:
																														_lowerCAmelCase  						=			loss_fct(logits.squeeze()   ,       labels.squeeze()				)
																								else:
																														_lowerCAmelCase  						=			loss_fct(_lowercase   ,       _lowercase				)
																		elif self.config.problem_type == "single_label_classification":
																								_lowerCAmelCase  						=			CrossEntropyLoss()
																								_lowerCAmelCase  						=			loss_fct(logits.view(-1   ,       self.num_labels				)   ,       labels.view(-1				)				)
																		elif self.config.problem_type == "multi_label_classification":
																								_lowerCAmelCase  						=			BCEWithLogitsLoss()
																								_lowerCAmelCase  						=			loss_fct(_lowercase   ,       _lowercase				)

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

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






@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    '''	,							_SCREAMING_SNAKE_CASE	,							)
class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase				):

												"""simple docstring"""



												super().__init__(_lowercase				)
												super()._init_backbone(_lowercase				)

												_lowerCAmelCase  						=			[config.embedding_size] + config.hidden_sizes
												_lowerCAmelCase  						=			ResNetEmbeddings(_lowercase				)
												_lowerCAmelCase  						=			ResNetEncoder(_lowercase				)

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




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

												"""simple docstring"""



												_lowerCAmelCase  						=			return_dict if return_dict is not None else self.config.use_return_dict
												_lowerCAmelCase  						=			(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)

												_lowerCAmelCase  						=			self.embedder(_lowercase				)

												_lowerCAmelCase  						=			self.encoder(_lowercase   ,       output_hidden_states=_lowercase   ,       return_dict=_lowercase				)

												_lowerCAmelCase  						=			outputs.hidden_states

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

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

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


'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


_lowercase    =		{"""configuration_deit""": ["""DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DeiTConfig""", """DeiTOnnxConfig"""]}

try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowercase    =		["""DeiTFeatureExtractor"""]
						_lowercase    =		["""DeiTImageProcessor"""]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowercase    =		[
						    """DEIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """DeiTForImageClassification""",
						    """DeiTForImageClassificationWithTeacher""",
						    """DeiTForMaskedImageModeling""",
						    """DeiTModel""",
						    """DeiTPreTrainedModel""",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowercase    =		[
						    """TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """TFDeiTForImageClassification""",
						    """TFDeiTForImageClassificationWithTeacher""",
						    """TFDeiTForMaskedImageModeling""",
						    """TFDeiTModel""",
						    """TFDeiTPreTrainedModel""",
						]


if TYPE_CHECKING:
						from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .feature_extraction_deit import DeiTFeatureExtractor
												from .image_processing_deit import DeiTImageProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_deit import (
												    DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    DeiTForImageClassification,
												    DeiTForImageClassificationWithTeacher,
												    DeiTForMaskedImageModeling,
												    DeiTModel,
												    DeiTPreTrainedModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_deit import (
												    TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFDeiTForImageClassification,
												    TFDeiTForImageClassificationWithTeacher,
												    TFDeiTForMaskedImageModeling,
												    TFDeiTModel,
												    TFDeiTPreTrainedModel,
												)


else:
						import sys

						_lowercase    =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
