

"""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, 1088, 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      __lowercase			(    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__					:  int					   =					nn.Convad(
       A					,A					,kernel_size=A					,stride=A					,padding=kernel_size // 2					,groups=A					,bias=A					,)
   UpperCAmelCase__					:  Optional[Any]					   =					nn.BatchNormad(A    )
   UpperCAmelCase__					:  List[Any]					   =					ACTaFN[activation] if activation is not None else nn.Identity()




  def    __lowercase       (  self		:			Dict					,A		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.convolution(A    )
   UpperCAmelCase__					:  Dict					   =					self.normalization(A    )
   UpperCAmelCase__					:  str					   =					self.activation(A    )
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			Optional[Any]					,A		:			RegNetConfig    ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__					:  str					   =					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__					:  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."""    )
   UpperCAmelCase__					:  Any					   =					self.embedder(A    )
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			Dict					,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__					:  List[str]					   =					nn.BatchNormad(A    )




  def    __lowercase       (  self		:			Any					,A		:			Tensor    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					self.convolution(A    )
   UpperCAmelCase__					:  List[Any]					   =					self.normalization(A    )
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			Tuple					,A		:			int					,A		:			int    ):






   '''simple docstring'''
   super().__init__()

   UpperCAmelCase__					:  List[str]					   =					nn.AdaptiveAvgPoolad((1, 1)    )
   UpperCAmelCase__					:  int					   =					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		:			Union[str, Any]    ):






   '''simple docstring'''
   # b c h w -> b c 1 1
   UpperCAmelCase__					:  Union[str, Any]					   =					self.pooler(A    )
   UpperCAmelCase__					:  Any					   =					self.attention(A    )
   UpperCAmelCase__					:  List[str]					   =					hidden_state * attention
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			List[str]					,A		:			RegNetConfig					,A		:			int					,A		:			int					,A		:			int = 1    ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__					:  int					   =					in_channels != out_channels or stride != 1
   UpperCAmelCase__					:  Optional[Any]					   =					max(1					,out_channels // config.groups_width    )
   UpperCAmelCase__					:  List[Any]					   =					(
       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__					:  Union[str, Any]					   =					ACTaFN[config.hidden_act]




  def    __lowercase       (  self		:			List[Any]					,A		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					hidden_state
   UpperCAmelCase__					:  Tuple					   =					self.layer(A    )
   UpperCAmelCase__					:  List[str]					   =					self.shortcut(A    )
   hidden_state += residual
   UpperCAmelCase__					:  Tuple					   =					self.activation(A    )
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			Optional[Any]					,A		:			RegNetConfig					,A		:			int					,A		:			int					,A		:			int = 1    ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__					:  Tuple					   =					in_channels != out_channels or stride != 1
   UpperCAmelCase__					:  Union[str, Any]					   =					max(1					,out_channels // config.groups_width    )
   UpperCAmelCase__					:  str					   =					(
       RegNetShortCut(A					,A					,stride=A    ) if should_apply_shortcut else nn.Identity()
   )
   UpperCAmelCase__					:  Optional[int]					   =					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__					:  Any					   =					ACTaFN[config.hidden_act]




  def    __lowercase       (  self		:			Any					,A		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					hidden_state
   UpperCAmelCase__					:  str					   =					self.layer(A    )
   UpperCAmelCase__					:  int					   =					self.shortcut(A    )
   hidden_state += residual
   UpperCAmelCase__					:  List[Any]					   =					self.activation(A    )
   return hidden_state







class      __lowercase			(    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__					:  Optional[Any]					   =					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		:			Dict					,A		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					self.layers(A    )
   return hidden_state







class      __lowercase			(    nn.Module					):

  def __init__(  self		:			int					,A		:			RegNetConfig    ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__					:  str					   =					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__					:  Dict					   =					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		:			List[str]					,A		:			Tensor					,A		:			bool = False					,A		:			bool = True    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					() if output_hidden_states else None

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

    UpperCAmelCase__					:  Dict					   =					stage_module(A    )

   if output_hidden_states:
    UpperCAmelCase__					:  Tuple					   =					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      __lowercase			(    __lowerCamelCase					):
  snake_case_				       =       RegNetConfig
  snake_case_				       =       """regnet"""
  snake_case_				       =       """pixel_values"""
  snake_case_				       =       True

  def    __lowercase       (  self		:			int					,A		:			Optional[int]    ):






   '''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		:			List[str]					,A		:			Any					,A		:			List[Any]=False    ):






   '''simple docstring'''
   if isinstance(A					,A    ):
    UpperCAmelCase__					:  Dict					   =					value







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

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







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

  def __init__(  self		:			Union[str, Any]					,A		:			Dict    ):






   '''simple docstring'''
   super().__init__(A    )
   UpperCAmelCase__					:  Tuple					   =					config
   UpperCAmelCase__					:  List[Any]					   =					RegNetEmbeddings(A    )
   UpperCAmelCase__					:  Optional[Any]					   =					RegNetEncoder(A    )
   UpperCAmelCase__					:  List[Any]					   =					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		:			Optional[Any]					,A		:			Tensor					,A		:			Optional[bool] = None					,A		:			Optional[bool] = None    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					(
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   UpperCAmelCase__					:  str					   =					return_dict if return_dict is not None else self.config.use_return_dict

   UpperCAmelCase__					:  Union[str, Any]					   =					self.embedder(A    )

   UpperCAmelCase__					:  int					   =					self.encoder(
       A					,output_hidden_states=A					,return_dict=A    )

   UpperCAmelCase__					:  Union[str, Any]					   =					encoder_outputs[0]

   UpperCAmelCase__					:  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(
    """
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """	,			__lowerCamelCase	,			)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class      __lowercase			(    __lowerCamelCase					):

  def __init__(  self		:			Optional[int]					,A		:			List[Any]    ):






   '''simple docstring'''
   super().__init__(A    )
   UpperCAmelCase__					:  Tuple					   =					config.num_labels
   UpperCAmelCase__					:  Dict					   =					RegNetModel(A    )
   # classification head
   UpperCAmelCase__					:  Optional[int]					   =					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		:			Dict					,A		:			Optional[torch.FloatTensor] = None					,A		:			Optional[torch.LongTensor] = None					,A		:			Optional[bool] = None					,A		:			Optional[bool] = None					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					return_dict if return_dict is not None else self.config.use_return_dict

   UpperCAmelCase__					:  Optional[Any]					   =					self.regnet(A					,output_hidden_states=A					,return_dict=A    )

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

   UpperCAmelCase__					:  Union[str, Any]					   =					self.classifier(A    )

   UpperCAmelCase__					:  str					   =					None

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

   if not return_dict:
    UpperCAmelCase__					:  Dict					   =					(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"""




# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__UpperCAmelCase =				abspath(join(dirname(dirname(__file__)), 'src'))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='ignore', category=FutureWarning)
def   lowerCAmelCase						(					__UpperCamelCase      ):
 '''simple docstring'''





 from diffusers.utils.testing_utils import pytest_addoption_shared

 pytest_addoption_shared(__UpperCamelCase      )




def   lowerCAmelCase						(					__UpperCamelCase      ):
 '''simple docstring'''





 from diffusers.utils.testing_utils import pytest_terminal_summary_main

 UpperCAmelCase__					:  Optional[Any]					   =					terminalreporter.config.getoption("""--make-reports"""      )
 if make_reports:
  pytest_terminal_summary_main(__UpperCamelCase							,    id=__UpperCamelCase      )





