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


__a: Union[str, Any]          =   logging.get_logger(__name__)

# General docstring
__a: Optional[int]          =   """RegNetConfig"""

# Base docstring
__a: List[Any]          =   """facebook/regnet-y-040"""
__a: List[str]          =   [1, 10_88, 7, 7]

# Image classification docstring
__a: Union[str, Any]          =   """facebook/regnet-y-040"""
__a: Optional[Any]          =   """tabby, tabby cat"""

__a: Optional[int]          =   [
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = 3	,   __lowerCAmelCase = 1	,   __lowerCAmelCase = 1	,   __lowerCAmelCase = "relu"	,   )  ->  Optional[int]:
        super().__init__()
        lowercase__				:    str  =					nn.Convad(
            __lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=__lowerCAmelCase	,   stride=__lowerCAmelCase	,   padding=kernel_size // 2	,   groups=__lowerCAmelCase	,   bias=__lowerCAmelCase	,   )
        lowercase__				:    str  =					nn.BatchNormad(__lowerCAmelCase )
        lowercase__				:    List[str]  =					ACTaFN[activation] if activation is not None else nn.Identity()


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  Any:
        lowercase__				:    Union[str, Any]  =					self.convolution(__lowerCAmelCase )
        lowercase__				:    Union[str, Any]  =					self.normalization(__lowerCAmelCase )
        lowercase__				:    Optional[int]  =					self.activation(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase )  ->  List[str]:
        super().__init__()
        lowercase__				:    Dict  =					RegNetConvLayer(
            config.num_channels	,   config.embedding_size	,   kernel_size=3	,   stride=2	,   activation=config.hidden_act )
        lowercase__				:    int  =					config.num_channels


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  List[str]:
        lowercase__				:    str  =					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.''' )
        lowercase__				:    Tuple  =					self.embedder(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = 2 )  ->  Optional[Any]:
        super().__init__()
        lowercase__				:    Union[str, Any]  =					nn.Convad(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1	,   stride=__lowerCAmelCase	,   bias=__lowerCAmelCase )
        lowercase__				:    Optional[int]  =					nn.BatchNormad(__lowerCAmelCase )


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  Tensor:
        lowercase__				:    Tuple  =					self.convolution(__lowerCAmelCase )
        lowercase__				:    Dict  =					self.normalization(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  List[Any]:
        super().__init__()

        lowercase__				:    Optional[int]  =					nn.AdaptiveAvgPoolad((1, 1) )
        lowercase__				:    str  =					nn.Sequential(
            nn.Convad(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1 )	,   nn.ReLU()	,   nn.Convad(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1 )	,   nn.Sigmoid()	,   )


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  Tuple:
        # b c h w -> b c 1 1
        lowercase__				:    List[str]  =					self.pooler(__lowerCAmelCase )
        lowercase__				:    Union[str, Any]  =					self.attention(__lowerCAmelCase )
        lowercase__				:    Dict  =					hidden_state * attention
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = 1 )  ->  Union[str, Any]:
        super().__init__()
        lowercase__				:    List[Any]  =					in_channels != out_channels or stride != 1
        lowercase__				:    List[Any]  =					max(1	,   out_channels // config.groups_width )
        lowercase__				:    Dict  =					(
            RegNetShortCut(__lowerCAmelCase	,   __lowerCAmelCase	,   stride=__lowerCAmelCase ) if should_apply_shortcut else nn.Identity()
        )
        lowercase__				:    Tuple  =					nn.Sequential(
            RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1	,   activation=config.hidden_act )	,   RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   stride=__lowerCAmelCase	,   groups=__lowerCAmelCase	,   activation=config.hidden_act )	,   RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1	,   activation=__lowerCAmelCase )	,   )
        lowercase__				:    Tuple  =					ACTaFN[config.hidden_act]


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  Tuple:
        lowercase__				:    Dict  =					hidden_state
        lowercase__				:    Tuple  =					self.layer(__lowerCAmelCase )
        lowercase__				:    int  =					self.shortcut(__lowerCAmelCase )
        hidden_state += residual
        lowercase__				:    Dict  =					self.activation(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = 1 )  ->  str:
        super().__init__()
        lowercase__				:    List[str]  =					in_channels != out_channels or stride != 1
        lowercase__				:    List[Any]  =					max(1	,   out_channels // config.groups_width )
        lowercase__				:    Optional[Any]  =					(
            RegNetShortCut(__lowerCAmelCase	,   __lowerCAmelCase	,   stride=__lowerCAmelCase ) if should_apply_shortcut else nn.Identity()
        )
        lowercase__				:    str  =					nn.Sequential(
            RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1	,   activation=config.hidden_act )	,   RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   stride=__lowerCAmelCase	,   groups=__lowerCAmelCase	,   activation=config.hidden_act )	,   RegNetSELayer(__lowerCAmelCase	,   reduced_channels=int(round(in_channels / 4 ) ) )	,   RegNetConvLayer(__lowerCAmelCase	,   __lowerCAmelCase	,   kernel_size=1	,   activation=__lowerCAmelCase )	,   )
        lowercase__				:    str  =					ACTaFN[config.hidden_act]


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  Dict:
        lowercase__				:    Tuple  =					hidden_state
        lowercase__				:    List[str]  =					self.layer(__lowerCAmelCase )
        lowercase__				:    List[Any]  =					self.shortcut(__lowerCAmelCase )
        hidden_state += residual
        lowercase__				:    Union[str, Any]  =					self.activation(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = 2	,   __lowerCAmelCase = 2	,   )  ->  int:
        super().__init__()

        lowercase__				:    Dict  =					RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

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


       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  str:
        lowercase__				:    List[str]  =					self.layers(__lowerCAmelCase )
        return hidden_state



class 			UpperCAmelCase  (						nn.Module					):





       '''simple docstring'''



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


       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase = False	,   __lowerCAmelCase = True )  ->  BaseModelOutputWithNoAttention:
        lowercase__				:    List[str]  =					() if output_hidden_states else None

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

         lowercase__				:    Optional[int]  =					stage_module(__lowerCAmelCase )

        if output_hidden_states:
         lowercase__				:    Optional[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=__lowerCAmelCase	,   hidden_states=__lowerCAmelCase )



class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE							=    RegNetConfig
       SCREAMING_SNAKE_CASE							=    "regnet"
       SCREAMING_SNAKE_CASE							=    "pixel_values"
       SCREAMING_SNAKE_CASE							=    True

       def _lowerCAmelCase(						self	,   __lowerCAmelCase )  ->  List[str]:
        if isinstance(__lowerCAmelCase	,   nn.Convad ):
         nn.init.kaiming_normal_(module.weight	,   mode='''fan_out'''	,   nonlinearity='''relu''' )
        elif isinstance(__lowerCAmelCase	,   (nn.BatchNormad, nn.GroupNorm) ):
         nn.init.constant_(module.weight	,   1 )
         nn.init.constant_(module.bias	,   0 )


       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase=False )  ->  str:
        if isinstance(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    List[Any]  =					value


__a: Tuple          =   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.
"""

__a: Tuple          =   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 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase )  ->  List[Any]:
        super().__init__(__lowerCAmelCase )
        lowercase__				:    List[Any]  =					config
        lowercase__				:    Tuple  =					RegNetEmbeddings(__lowerCAmelCase )
        lowercase__				:    List[Any]  =					RegNetEncoder(__lowerCAmelCase )
        lowercase__				:    int  =					nn.AdaptiveAvgPoolad((1, 1) )
        # Initialize weights and apply final processing
        self.post_init()


       @add_start_docstrings_to_model_forward(__lowerCAmelCase )
       @add_code_sample_docstrings(
           checkpoint=_CHECKPOINT_FOR_DOC	,   output_type=__lowerCAmelCase	,   config_class=_CONFIG_FOR_DOC	,   modality='''vision'''	,   expected_output=_EXPECTED_OUTPUT_SHAPE	,   )
       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None )  ->  BaseModelOutputWithPoolingAndNoAttention:
        lowercase__				:    Union[str, Any]  =					(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowercase__				:    Dict  =					return_dict if return_dict is not None else self.config.use_return_dict

        lowercase__				:    int  =					self.embedder(__lowerCAmelCase )

        lowercase__				:    List[Any]  =					self.encoder(
            __lowerCAmelCase	,   output_hidden_states=__lowerCAmelCase	,   return_dict=__lowerCAmelCase )

        lowercase__				:    List[Any]  =					encoder_outputs[0]

        lowercase__				:    Union[str, Any]  =					self.pooler(__lowerCAmelCase )

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

        return BaseModelOutputWithPoolingAndNoAttention(
            last_hidden_state=__lowerCAmelCase	,   pooler_output=__lowerCAmelCase	,   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 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase )  ->  str:
        super().__init__(__lowerCAmelCase )
        lowercase__				:    List[Any]  =					config.num_labels
        lowercase__				:    Union[str, Any]  =					RegNetModel(__lowerCAmelCase )
        # classification head
        lowercase__				:    List[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(__lowerCAmelCase )
       @add_code_sample_docstrings(
           checkpoint=_IMAGE_CLASS_CHECKPOINT	,   output_type=__lowerCAmelCase	,   config_class=_CONFIG_FOR_DOC	,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,   )
       def _lowerCAmelCase(						self	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   )  ->  ImageClassifierOutputWithNoAttention:
        lowercase__				:    Tuple  =					return_dict if return_dict is not None else self.config.use_return_dict

        lowercase__				:    int  =					self.regnet(__lowerCAmelCase	,   output_hidden_states=__lowerCAmelCase	,   return_dict=__lowerCAmelCase )

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

        lowercase__				:    Optional[Any]  =					self.classifier(__lowerCAmelCase )

        lowercase__				:    Union[str, Any]  =					None

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

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

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


'''simple docstring'''

import copy

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


__a: List[str]          =   logging.get_logger(__name__)



class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE							=    "encoder-decoder"
       SCREAMING_SNAKE_CASE							=    True

       def __init__(						self	,   **__lowerCAmelCase )  ->  int:
        super().__init__(**__lowerCAmelCase )
        assert (
            "encoder" in kwargs and "decoder" in kwargs
        ), "Config has to be initialized with encoder and decoder config"
        lowercase__				:    Optional[int]  =					kwargs.pop('''encoder''' )
        lowercase__				:    Union[str, Any]  =					encoder_config.pop('''model_type''' )
        lowercase__				:    Any  =					kwargs.pop('''decoder''' )
        lowercase__				:    Any  =					decoder_config.pop('''model_type''' )

        from ..auto.configuration_auto import AutoConfig

        lowercase__				:    Union[str, Any]  =					AutoConfig.for_model(__lowerCAmelCase	,   **__lowerCAmelCase )
        lowercase__				:    Optional[Any]  =					AutoConfig.for_model(__lowerCAmelCase	,   **__lowerCAmelCase )
        lowercase__				:    Tuple  =					True

       @classmethod
       def _lowerCAmelCase(						cls	,   __lowerCAmelCase	,   __lowerCAmelCase	,   **__lowerCAmelCase )  ->  PretrainedConfig:
        logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config''' )
        lowercase__				:    Union[str, Any]  =					True
        lowercase__				:    Any  =					True

        return cls(encoder=encoder_config.to_dict()	,   decoder=decoder_config.to_dict()	,   **__lowerCAmelCase )


       def _lowerCAmelCase(						self )  ->  Any:
        lowercase__				:    Any  =					copy.deepcopy(self.__dict__ )
        lowercase__				:    Optional[Any]  =					self.encoder.to_dict()
        lowercase__				:    Tuple  =					self.decoder.to_dict()
        lowercase__				:    Dict  =					self.__class__.model_type
        return output

