


'''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 ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCamelCase__          = logging.get_logger(__name__)

# General docstring
UpperCamelCase__          = '''RegNetConfig'''

# Base docstring
UpperCamelCase__          = '''facebook/regnet-y-040'''
UpperCamelCase__          = [1, 1_0_8_8, 7, 7]

# Image classification docstring
UpperCamelCase__          = '''facebook/regnet-y-040'''
UpperCamelCase__          = '''tabby, tabby cat'''

UpperCamelCase__          = [
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       Optional[Any]		,		_A							:       int		,		_A							:       int		,		_A							:       int = 3		,		_A							:       int = 1		,		_A							:       int = 1		,		_A							:       Optional[str] = "relu"		,		):


          '''simple docstring'''


          super().__init__()
          UpperCAmelCase__     :						Optional[Any]						=							nn.Convad(
              _A		,		_A		,		kernel_size=_A		,		stride=_A		,		padding=kernel_size // 2		,		groups=_A		,		bias=_A		,		)
          UpperCAmelCase__     :						str						=							nn.BatchNormad(_A					)
          UpperCAmelCase__     :						str						=							ACTaFN[activation] if activation is not None else nn.Identity()






      def lowercase_				(   self							:       str		,		_A							:       List[Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						str						=							self.convolution(_A					)
          UpperCAmelCase__     :						str						=							self.normalization(_A					)
          UpperCAmelCase__     :						Optional[int]						=							self.activation(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       Tuple		,		_A							:       RegNetConfig					):


          '''simple docstring'''


          super().__init__()
          UpperCAmelCase__     :						Optional[Any]						=							RegNetConvLayer(
              config.num_channels		,		config.embedding_size		,		kernel_size=3		,		stride=2		,		activation=config.hidden_act					)
          UpperCAmelCase__     :						List[str]						=							config.num_channels






      def lowercase_				(   self							:       Dict		,		_A							:       str					):


          '''simple docstring'''


          UpperCAmelCase__     :						Tuple						=							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.'''					)
          UpperCAmelCase__     :						Tuple						=							self.embedder(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       str		,		_A							:       int		,		_A							:       int		,		_A							:       int = 2					):


          '''simple docstring'''


          super().__init__()
          UpperCAmelCase__     :						List[str]						=							nn.Convad(_A		,		_A		,		kernel_size=1		,		stride=_A		,		bias=_A					)
          UpperCAmelCase__     :						Optional[int]						=							nn.BatchNormad(_A					)






      def lowercase_				(   self							:       int		,		_A							:       Tensor					):


          '''simple docstring'''


          UpperCAmelCase__     :						int						=							self.convolution(_A					)
          UpperCAmelCase__     :						List[Any]						=							self.normalization(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       List[str]		,		_A							:       int		,		_A							:       int					):


          '''simple docstring'''


          super().__init__()

          UpperCAmelCase__     :						List[str]						=							nn.AdaptiveAvgPoolad((1, 1)					)
          UpperCAmelCase__     :						Dict						=							nn.Sequential(
              nn.Convad(_A		,		_A		,		kernel_size=1					)		,		nn.ReLU()		,		nn.Convad(_A		,		_A		,		kernel_size=1					)		,		nn.Sigmoid()		,		)






      def lowercase_				(   self							:       List[str]		,		_A							:       str					):


          '''simple docstring'''


          UpperCAmelCase__     :						Any						=							self.pooler(_A					)
          UpperCAmelCase__     :						Optional[Any]						=							self.attention(_A					)
          UpperCAmelCase__     :						int						=							hidden_state * attention
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       int		,		_A							:       RegNetConfig		,		_A							:       int		,		_A							:       int		,		_A							:       int = 1					):


          '''simple docstring'''


          super().__init__()
          UpperCAmelCase__     :						str						=							in_channels != out_channels or stride != 1
          UpperCAmelCase__     :						Optional[Any]						=							max(1		,		out_channels // config.groups_width					)
          UpperCAmelCase__     :						List[str]						=							(
              RegNetShortCut(_A		,		_A		,		stride=_A					) if should_apply_shortcut else nn.Identity()
          )
          UpperCAmelCase__     :						Union[str, Any]						=							nn.Sequential(
              RegNetConvLayer(_A		,		_A		,		kernel_size=1		,		activation=config.hidden_act					)		,		RegNetConvLayer(_A		,		_A		,		stride=_A		,		groups=_A		,		activation=config.hidden_act					)		,		RegNetConvLayer(_A		,		_A		,		kernel_size=1		,		activation=_A					)		,		)
          UpperCAmelCase__     :						Optional[int]						=							ACTaFN[config.hidden_act]






      def lowercase_				(   self							:       List[str]		,		_A							:       Union[str, Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						int						=							hidden_state
          UpperCAmelCase__     :						List[Any]						=							self.layer(_A					)
          UpperCAmelCase__     :						List[Any]						=							self.shortcut(_A					)
          hidden_state += residual
          UpperCAmelCase__     :						str						=							self.activation(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       List[Any]		,		_A							:       RegNetConfig		,		_A							:       int		,		_A							:       int		,		_A							:       int = 1					):


          '''simple docstring'''


          super().__init__()
          UpperCAmelCase__     :						Optional[int]						=							in_channels != out_channels or stride != 1
          UpperCAmelCase__     :						Dict						=							max(1		,		out_channels // config.groups_width					)
          UpperCAmelCase__     :						Union[str, Any]						=							(
              RegNetShortCut(_A		,		_A		,		stride=_A					) if should_apply_shortcut else nn.Identity()
          )
          UpperCAmelCase__     :						List[str]						=							nn.Sequential(
              RegNetConvLayer(_A		,		_A		,		kernel_size=1		,		activation=config.hidden_act					)		,		RegNetConvLayer(_A		,		_A		,		stride=_A		,		groups=_A		,		activation=config.hidden_act					)		,		RegNetSELayer(_A		,		reduced_channels=int(round(in_channels / 4					)					)					)		,		RegNetConvLayer(_A		,		_A		,		kernel_size=1		,		activation=_A					)		,		)
          UpperCAmelCase__     :						Tuple						=							ACTaFN[config.hidden_act]






      def lowercase_				(   self							:       Dict		,		_A							:       Union[str, Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						List[str]						=							hidden_state
          UpperCAmelCase__     :						List[str]						=							self.layer(_A					)
          UpperCAmelCase__     :						Tuple						=							self.shortcut(_A					)
          hidden_state += residual
          UpperCAmelCase__     :						Union[str, Any]						=							self.activation(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       Optional[Any]		,		_A							:       RegNetConfig		,		_A							:       int		,		_A							:       int		,		_A							:       int = 2		,		_A							:       int = 2		,		):


          '''simple docstring'''


          super().__init__()

          UpperCAmelCase__     :						Union[str, Any]						=							RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

          UpperCAmelCase__     :						str						=							nn.Sequential(
              # downsampling is done in the first layer with stride of 2
              layer(
                  _A		,		_A		,		_A		,		stride=_A		,		)		,		*[layer(_A		,		_A		,		_A					) for _ in range(depth - 1					)]		,		)






      def lowercase_				(   self							:       Tuple		,		_A							:       Optional[Any]					):


          '''simple docstring'''


          UpperCAmelCase__     :						Union[str, Any]						=							self.layers(_A					)
          return hidden_state




class      lowerCamelCase_  (      nn.Module ):






      def __init__(   self							:       str		,		_A							:       RegNetConfig					):


          '''simple docstring'''


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






      def lowercase_				(   self							:       Optional[int]		,		_A							:       Tensor		,		_A							:       bool = False		,		_A							:       bool = True					):


          '''simple docstring'''


          UpperCAmelCase__     :						Union[str, Any]						=							() if output_hidden_states else None

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

              UpperCAmelCase__     :						str						=							stage_module(_A					)

          if output_hidden_states:
              UpperCAmelCase__     :						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      lowerCamelCase_  (      __a ):
      lowerCAmelCase__							=	RegNetConfig
      lowerCAmelCase__							=	'regnet'
      lowerCAmelCase__							=	'pixel_values'
      lowerCAmelCase__							=	True






      def lowercase_				(   self							:       Optional[Any]		,		_A							:       Optional[Any]					):


          '''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 lowercase_				(   self							:       Union[str, Any]		,		_A							:       int		,		_A							:       Dict=False					):


          '''simple docstring'''


          if isinstance(_A		,		_A					):
              UpperCAmelCase__     :						Dict						=							value

UpperCamelCase__          = 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 ([`RegNetConfig`]): 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.
'''

UpperCamelCase__          = 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 [`~file_utils.ModelOutput`] instead of a plain tuple.
'''




@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.'	,			__a	,			)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class      lowerCamelCase_  (      __a ):






      def __init__(   self							:       Union[str, Any]		,		_A							:       Tuple					):


          '''simple docstring'''


          super().__init__(_A					)
          UpperCAmelCase__     :						Union[str, Any]						=							config
          UpperCAmelCase__     :						Optional[int]						=							RegNetEmbeddings(_A					)
          UpperCAmelCase__     :						Dict						=							RegNetEncoder(_A					)
          UpperCAmelCase__     :						Dict						=							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 lowercase_				(   self							:       List[Any]		,		_A							:       Tensor		,		_A							:       Optional[bool] = None		,		_A							:       Optional[bool] = None					):


          '''simple docstring'''


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

          UpperCAmelCase__     :						str						=							self.embedder(_A					)

          UpperCAmelCase__     :						Optional[Any]						=							self.encoder(
              _A		,		output_hidden_states=_A		,		return_dict=_A					)

          UpperCAmelCase__     :						str						=							encoder_outputs[0]

          UpperCAmelCase__     :						Union[str, Any]						=							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(
    '\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '	,			__a	,			)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class      lowerCamelCase_  (      __a ):






      def __init__(   self							:       Any		,		_A							:       Tuple					):


          '''simple docstring'''


          super().__init__(_A					)
          UpperCAmelCase__     :						Union[str, Any]						=							config.num_labels
          UpperCAmelCase__     :						Union[str, Any]						=							RegNetModel(_A					)
          # classification head
          UpperCAmelCase__     :						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 lowercase_				(   self							:       List[str]		,		_A							:       Optional[torch.FloatTensor] = None		,		_A							:       Optional[torch.LongTensor] = None		,		_A							:       Optional[bool] = None		,		_A							:       Optional[bool] = None		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						List[Any]						=							return_dict if return_dict is not None else self.config.use_return_dict

          UpperCAmelCase__     :						Union[str, Any]						=							self.regnet(_A		,		output_hidden_states=_A		,		return_dict=_A					)

          UpperCAmelCase__     :						Tuple						=							outputs.pooler_output if return_dict else outputs[1]

          UpperCAmelCase__     :						Dict						=							self.classifier(_A					)

          UpperCAmelCase__     :						str						=							None

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

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

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



'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_a					:			Any							=    {
    'configuration_groupvit': [
        'GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'GroupViTConfig',
        'GroupViTOnnxConfig',
        'GroupViTTextConfig',
        'GroupViTVisionConfig',
    ],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _a					:			Union[str, Any]							=    [
      'GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'GroupViTModel',
      'GroupViTPreTrainedModel',
      'GroupViTTextModel',
      'GroupViTVisionModel',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _a					:			Tuple							=    [
      'TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'TFGroupViTModel',
      'TFGroupViTPreTrainedModel',
      'TFGroupViTTextModel',
      'TFGroupViTVisionModel',
  ]

if TYPE_CHECKING:
  from .configuration_groupvit import (
      GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      GroupViTConfig,
      GroupViTOnnxConfig,
      GroupViTTextConfig,
      GroupViTVisionConfig,
  )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_groupvit import (
        GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        GroupViTModel,
        GroupViTPreTrainedModel,
        GroupViTTextModel,
        GroupViTVisionModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_groupvit import (
        TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFGroupViTModel,
        TFGroupViTPreTrainedModel,
        TFGroupViTTextModel,
        TFGroupViTVisionModel,
    )

else:
  import sys

  _a					:			Union[str, Any]							=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


