






"""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


__lowerCamelCase :      str										=				logging.get_logger(__name__)

# General docstring
__lowerCamelCase :      Union[str, Any]										=				'ResNetConfig'

# Base docstring
__lowerCamelCase :      Tuple										=				'microsoft/resnet-50'
__lowerCamelCase :      str										=				[1, 2048, 7, 7]

# Image classification docstring
__lowerCamelCase :      Any										=				'microsoft/resnet-50'
__lowerCamelCase :      Tuple										=				'tiger cat'

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



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  int							,		__A					:  int							,		__A					:  int							,		__A					:  int = 3							,		__A					:  int = 1							,		__A					:  str = "relu"	):
						super().__init__()
						snake_case__			:				List[Any]       =  nn.Convad(
						    _lowercase							,		_lowercase							,		kernel_size=_lowercase							,		stride=_lowercase							,		padding=kernel_size // 2							,		bias=_lowercase	)
						snake_case__			:				Dict       =  nn.BatchNormad(_lowercase	)
						snake_case__			:				Dict       =  ACTaFN[activation] if activation is not None else nn.Identity()






					def 	_lowercase					( self					:  str							,		__A					:  Tensor	):
						snake_case__			:				Optional[int]       =  self.convolution(_lowercase	)
						snake_case__			:				str       =  self.normalization(_lowercase	)
						snake_case__			:				Optional[Any]       =  self.activation(_lowercase	)
						return hidden_state



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  List[str]							,		__A					:  ResNetConfig	):
						super().__init__()
						snake_case__			:				str       =  ResNetConvLayer(
						    config.num_channels							,		config.embedding_size							,		kernel_size=7							,		stride=2							,		activation=config.hidden_act	)
						snake_case__			:				List[Any]       =  nn.MaxPoolad(kernel_size=3							,		stride=2							,		padding=1	)
						snake_case__			:				str       =  config.num_channels






					def 	_lowercase					( self					:  Optional[int]							,		__A					:  Tensor	):
						snake_case__			:				Union[str, Any]       =  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."	)
						snake_case__			:				List[str]       =  self.embedder(_lowercase	)
						snake_case__			:				Dict       =  self.pooler(_lowercase	)
						return embedding



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  int							,		__A					:  int							,		__A					:  int = 2	):
						super().__init__()
						snake_case__			:				List[str]       =  nn.Convad(_lowercase							,		_lowercase							,		kernel_size=1							,		stride=_lowercase							,		bias=_lowercase	)
						snake_case__			:				Dict       =  nn.BatchNormad(_lowercase	)






					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  Tensor	):
						snake_case__			:				Optional[int]       =  self.convolution(_lowercase	)
						snake_case__			:				Dict       =  self.normalization(_lowercase	)
						return hidden_state



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  Any							,		__A					:  int							,		__A					:  int							,		__A					:  int = 1							,		__A					:  str = "relu"	):
						super().__init__()
						snake_case__			:				str       =  in_channels != out_channels or stride != 1
						snake_case__			:				List[str]       =  (
						    ResNetShortCut(_lowercase							,		_lowercase							,		stride=_lowercase	) if should_apply_shortcut else nn.Identity()
						)
						snake_case__			:				Dict       =  nn.Sequential(
						    ResNetConvLayer(_lowercase							,		_lowercase							,		stride=_lowercase	)							,		ResNetConvLayer(_lowercase							,		_lowercase							,		activation=_lowercase	)							,		)
						snake_case__			:				Tuple       =  ACTaFN[activation]






					def 	_lowercase					( self					:  str							,		__A					:  Optional[int]	):
						snake_case__			:				List[Any]       =  hidden_state
						snake_case__			:				int       =  self.layer(_lowercase	)
						snake_case__			:				Optional[int]       =  self.shortcut(_lowercase	)
						hidden_state += residual
						snake_case__			:				Any       =  self.activation(_lowercase	)
						return hidden_state



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  int							,		__A					:  int							,		__A					:  int = 1							,		__A					:  str = "relu"							,		__A					:  int = 4	):
						super().__init__()
						snake_case__			:				Any       =  in_channels != out_channels or stride != 1
						snake_case__			:				Tuple       =  out_channels // reduction
						snake_case__			:				Optional[Any]       =  (
						    ResNetShortCut(_lowercase							,		_lowercase							,		stride=_lowercase	) if should_apply_shortcut else nn.Identity()
						)
						snake_case__			:				List[Any]       =  nn.Sequential(
						    ResNetConvLayer(_lowercase							,		_lowercase							,		kernel_size=1	)							,		ResNetConvLayer(_lowercase							,		_lowercase							,		stride=_lowercase	)							,		ResNetConvLayer(_lowercase							,		_lowercase							,		kernel_size=1							,		activation=_lowercase	)							,		)
						snake_case__			:				Optional[Any]       =  ACTaFN[activation]






					def 	_lowercase					( self					:  Dict							,		__A					:  Union[str, Any]	):
						snake_case__			:				Dict       =  hidden_state
						snake_case__			:				List[str]       =  self.layer(_lowercase	)
						snake_case__			:				Dict       =  self.shortcut(_lowercase	)
						hidden_state += residual
						snake_case__			:				Optional[int]       =  self.activation(_lowercase	)
						return hidden_state



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  Union[str, Any]							,		__A					:  ResNetConfig							,		__A					:  int							,		__A					:  int							,		__A					:  int = 2							,		__A					:  int = 2							,		):
						super().__init__()

						snake_case__			:				str       =  ResNetBottleNeckLayer if config.layer_type == "bottleneck" else ResNetBasicLayer

						snake_case__			:				Tuple       =  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					:  List[str]							,		__A					:  Tensor	):
						snake_case__			:				Any       =  input
						for layer in self.layers:
							snake_case__			:				int       =  layer(_lowercase	)
						return hidden_state



class 							SCREAMING_SNAKE_CASE__						(    nn.Module						):





					"""simple docstring"""






					def __init__( self					:  Any							,		__A					:  ResNetConfig	):
						super().__init__()
						snake_case__			:				Dict       =  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]							,		)	)
						snake_case__			:				Optional[Any]       =  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					:  Optional[Any]							,		__A					:  Tensor							,		__A					:  bool = False							,		__A					:  bool = True	):
						snake_case__			:				Dict       =  () if output_hidden_states else None

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

							snake_case__			:				str       =  stage_module(_lowercase	)

						if output_hidden_states:
							snake_case__			:				int       =  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 							SCREAMING_SNAKE_CASE__						(    _lowerCAmelCase						):





					"""simple docstring"""






					a_     =		ResNetConfig
					a_     =		"resnet"
					a_     =		"pixel_values"
					a_     =		True



					def 	_lowercase					( self					:  Any							,		__A					:  Optional[int]	):
						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					:  Optional[Any]							,		__A					:  Dict							,		__A					:  int=False	):
						if isinstance(_lowercase							,		_lowercase	):
							snake_case__			:				List[Any]       =  value
__lowerCamelCase :      List[Any]										=				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 :      List[Any]										=				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 							SCREAMING_SNAKE_CASE__						(    _lowerCAmelCase						):





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  Dict	):
						super().__init__(_lowercase	)
						snake_case__			:				Any       =  config
						snake_case__			:				int       =  ResNetEmbeddings(_lowercase	)
						snake_case__			:				Any       =  ResNetEncoder(_lowercase	)
						snake_case__			:				Optional[Any]       =  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					:  str							,		__A					:  Tensor							,		__A					:  Optional[bool] = None							,		__A					:  Optional[bool] = None	):
						snake_case__			:				Tuple       =  (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						snake_case__			:				Optional[int]       =  return_dict if return_dict is not None else self.config.use_return_dict

						snake_case__			:				int       =  self.embedder(_lowercase	)

						snake_case__			:				Union[str, Any]       =  self.encoder(
						    _lowercase							,		output_hidden_states=_lowercase							,		return_dict=_lowercase	)

						snake_case__			:				List[Any]       =  encoder_outputs[0]

						snake_case__			:				Optional[int]       =  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(
    "\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "			,       _lowerCAmelCase			,       )
class 							SCREAMING_SNAKE_CASE__						(    _lowerCAmelCase						):





					"""simple docstring"""






					def __init__( self					:  Union[str, Any]							,		__A					:  Tuple	):
						super().__init__(_lowercase	)
						snake_case__			:				List[str]       =  config.num_labels
						snake_case__			:				Optional[int]       =  ResNetModel(_lowercase	)
						# classification head
						snake_case__			:				Any       =  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					:  int							,		__A					:  Optional[torch.FloatTensor] = None							,		__A					:  Optional[torch.LongTensor] = None							,		__A					:  Optional[bool] = None							,		__A					:  Optional[bool] = None							,		):
						snake_case__			:				Optional[int]       =  return_dict if return_dict is not None else self.config.use_return_dict

						snake_case__			:				Optional[int]       =  self.resnet(_lowercase							,		output_hidden_states=_lowercase							,		return_dict=_lowercase	)

						snake_case__			:				Any       =  outputs.pooler_output if return_dict else outputs[1]

						snake_case__			:				Union[str, Any]       =  self.classifier(_lowercase	)

						snake_case__			:				Union[str, Any]       =  None

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

						if not return_dict:
							snake_case__			:				str       =  (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(
    "\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    "			,       _lowerCAmelCase			,       )
class 							SCREAMING_SNAKE_CASE__						(    _lowerCAmelCase			,       _lowerCAmelCase						):





					"""simple docstring"""






					def __init__( self					:  Optional[int]							,		__A					:  int	):
						super().__init__(_lowercase	)
						super()._init_backbone(_lowercase	)

						snake_case__			:				Dict       =  [config.embedding_size] + config.hidden_sizes
						snake_case__			:				Optional[Any]       =  ResNetEmbeddings(_lowercase	)
						snake_case__			:				Any       =  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					:  int							,		__A					:  Tensor							,		__A					:  Optional[bool] = None							,		__A					:  Optional[bool] = None	):
						snake_case__			:				List[str]       =  return_dict if return_dict is not None else self.config.use_return_dict
						snake_case__			:				str       =  (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)

						snake_case__			:				int       =  self.embedder(_lowercase	)

						snake_case__			:				Dict       =  self.encoder(_lowercase							,		output_hidden_states=_lowercase							,		return_dict=_lowercase	)

						snake_case__			:				Dict       =  outputs.hidden_states

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

						if not return_dict:
							snake_case__			:				Dict       =  (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							,		)











from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch



@require_torch
@require_tf
class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					def 	_lowercase					( self					:  Optional[Any]	):
						snake_case__			:				Optional[Any]       =  SMALL_MODEL_IDENTIFIER
						snake_case__			:				Any       =  "pt"
						snake_case__			:				Any       =  "tf"



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  List[Any]	):
						snake_case__			:				int       =  AutoModel.from_pretrained(self.test_model	)
						model_pt.save_pretrained(__A	)



					def 	_lowercase					( self					:  Optional[int]							,		__A					:  Tuple	):
						snake_case__			:				List[Any]       =  TFAutoModel.from_pretrained(self.test_model							,		from_pt=__A	)
						model_tf.save_pretrained(__A	)



					def 	_lowercase					( self					:  str	):
						snake_case__			:				Optional[Any]       =  "mock_framework"

						# Framework provided - return whatever the user provides
						snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model							,		__A	)
						self.assertEqual(__A							,		__A	)

						# Local checkpoint and framework provided - return provided framework
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(__A	)
							snake_case__			:				Optional[int]       =  FeaturesManager.determine_framework(__A							,		__A	)
							self.assertEqual(__A							,		__A	)

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(__A	)
							snake_case__			:				int       =  FeaturesManager.determine_framework(__A							,		__A	)
							self.assertEqual(__A							,		__A	)



					def 	_lowercase					( self					:  Dict	):
						# PyTorch checkpoint
						with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(__A	)
							snake_case__			:				List[str]       =  FeaturesManager.determine_framework(__A	)
							self.assertEqual(__A							,		self.framework_pt	)

						# TensorFlow checkpoint
						with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(__A	)
							snake_case__			:				Tuple       =  FeaturesManager.determine_framework(__A	)
							self.assertEqual(__A							,		self.framework_tf	)

						# Invalid local checkpoint
						with TemporaryDirectory() as local_invalid_ckpt:
							with self.assertRaises(__A	):
								snake_case__			:				int       =  FeaturesManager.determine_framework(__A	)






					def 	_lowercase					( self					:  Dict	):
						snake_case__			:				Dict       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	):
							snake_case__			:				List[Any]       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_pt	)

						# PyTorch not in environment -> use TensorFlow
						snake_case__			:				Tuple       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_torch_available"							,		__A	):
							snake_case__			:				int       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_tf	)

						# Both in environment -> use PyTorch
						snake_case__			:				Dict       =  MagicMock(return_value=__A	)
						snake_case__			:				Optional[int]       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	), patch(
						    "transformers.onnx.features.is_torch_available"							,		__A	):
							snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model	)
							self.assertEqual(__A							,		self.framework_pt	)

						# Both not in environment -> raise error
						snake_case__			:				List[str]       =  MagicMock(return_value=__A	)
						snake_case__			:				Optional[Any]       =  MagicMock(return_value=__A	)
						with patch("transformers.onnx.features.is_tf_available"							,		__A	), patch(
						    "transformers.onnx.features.is_torch_available"							,		__A	):
							with self.assertRaises(__A	):
								snake_case__			:				Optional[Any]       =  FeaturesManager.determine_framework(self.test_model	)



