





"""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_ =			logging.get_logger(__name__)

# General docstring
a_ =			"""RegNetConfig"""

# Base docstring
a_ =			"""facebook/regnet-y-040"""
a_ =			[1, 1088, 7, 7]

# Image classification docstring
a_ =			"""facebook/regnet-y-040"""
a_ =			"""tabby, tabby cat"""

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






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


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



     '''simple docstring'''




     super().__init__()
     __A    :						int						 =     nn.Convad(
         __lowerCamelCase					,						__lowerCamelCase					,						kernel_size=__lowerCamelCase					,						stride=__lowerCamelCase					,						padding=kernel_size // 2					,						groups=__lowerCamelCase					,						bias=__lowerCamelCase					,						)
     __A    :						List[str]						 =     nn.BatchNormad(__lowerCamelCase  )
     __A    :						List[str]						 =     ACTaFN[activation] if activation is not None else nn.Identity()







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     self.convolution(__lowerCamelCase  )
     __A    :						List[str]						 =     self.normalization(__lowerCamelCase  )
     __A    :						List[str]						 =     self.activation(__lowerCamelCase  )
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




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







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     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.'''  )
     __A    :						Optional[Any]						 =     self.embedder(__lowerCamelCase  )
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


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



     '''simple docstring'''




     super().__init__()
     __A    :						Tuple						 =     nn.Convad(__lowerCamelCase					,						__lowerCamelCase					,						kernel_size=1					,						stride=__lowerCamelCase					,						bias=__lowerCamelCase  )
     __A    :						Dict						 =     nn.BatchNormad(__lowerCamelCase  )







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     self.convolution(__lowerCamelCase  )
     __A    :						Union[str, Any]						 =     self.normalization(__lowerCamelCase  )
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     super().__init__()

     __A    :						Tuple						 =     nn.AdaptiveAvgPoolad((1, 1)  )
     __A    :						int						 =     nn.Sequential(
         nn.Convad(__lowerCamelCase					,						__lowerCamelCase					,						kernel_size=1  )					,						nn.ReLU()					,						nn.Convad(__lowerCamelCase					,						__lowerCamelCase					,						kernel_size=1  )					,						nn.Sigmoid()					,						)







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						int						 =     self.pooler(__lowerCamelCase  )
     __A    :						Union[str, Any]						 =     self.attention(__lowerCamelCase  )
     __A    :						int						 =     hidden_state * attention
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase = 1  ):



     '''simple docstring'''




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







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						str						 =     hidden_state
     __A    :						List[str]						 =     self.layer(__lowerCamelCase  )
     __A    :						Dict						 =     self.shortcut(__lowerCamelCase  )
     hidden_state += residual
     __A    :						List[Any]						 =     self.activation(__lowerCamelCase  )
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase = 1  ):



     '''simple docstring'''




     super().__init__()
     __A    :						Union[str, Any]						 =     in_channels != out_channels or stride != 1
     __A    :						str						 =     max(1					,						out_channels // config.groups_width  )
     __A    :						Dict						 =     (
         RegNetShortCut(__lowerCamelCase					,						__lowerCamelCase					,						stride=__lowerCamelCase  ) if should_apply_shortcut else nn.Identity()
     )
     __A    :						Any						 =     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  )					,						)
     __A    :						int						 =     ACTaFN[config.hidden_act]







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




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






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


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



     '''simple docstring'''




     super().__init__()

     __A    :						Tuple						 =     RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

     __A    :						Tuple						 =     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 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Any						 =     self.layers(__lowerCamelCase  )
     return hidden_state






class 		__snake_case					(		nn.Module     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     super().__init__()
     __A    :						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(
             __lowerCamelCase					,						config.embedding_size					,						config.hidden_sizes[0]					,						stride=2 if config.downsample_in_first_stage else 1					,						depth=config.depths[0]					,						)  )
     __A    :						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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase = False					,						__lowerCamelCase = True  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     () if output_hidden_states else None

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

         __A    :						str						 =     stage_module(__lowerCamelCase  )

     if output_hidden_states:
         __A    :						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 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     RegNetConfig
 _lowerCamelCase        =     """regnet"""
 _lowerCamelCase        =     """pixel_values"""
 _lowerCamelCase        =     True

 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase=False  ):



     '''simple docstring'''




     if isinstance(__lowerCamelCase					,						__lowerCamelCase  ):
         __A    :						List[Any]						 =     value






a_ =			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_ =			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."""							,							SCREAMING_SNAKE_CASE__							,							)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     super().__init__(__lowerCamelCase  )
     __A    :						Dict						 =     config
     __A    :						List[str]						 =     RegNetEmbeddings(__lowerCamelCase  )
     __A    :						Dict						 =     RegNetEncoder(__lowerCamelCase  )
     __A    :						List[str]						 =     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase = None					,						__lowerCamelCase = None  ):



     '''simple docstring'''




     __A    :						Tuple						 =     (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     __A    :						Tuple						 =     return_dict if return_dict is not None else self.config.use_return_dict

     __A    :						str						 =     self.embedder(__lowerCamelCase  )

     __A    :						List[Any]						 =     self.encoder(
         __lowerCamelCase					,						output_hidden_states=__lowerCamelCase					,						return_dict=__lowerCamelCase  )

     __A    :						Dict						 =     encoder_outputs[0]

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





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     super().__init__(__lowerCamelCase  )
     __A    :						Any						 =     config.num_labels
     __A    :						Dict						 =     RegNetModel(__lowerCamelCase  )
     # classification head
     __A    :						Dict						 =     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 					UpperCamelCase__(    self					,						__lowerCamelCase = None					,						__lowerCamelCase = None					,						__lowerCamelCase = None					,						__lowerCamelCase = None					,						):



     '''simple docstring'''




     __A    :						int						 =     return_dict if return_dict is not None else self.config.use_return_dict

     __A    :						Any						 =     self.regnet(__lowerCamelCase					,						output_hidden_states=__lowerCamelCase					,						return_dict=__lowerCamelCase  )

     __A    :						str						 =     outputs.pooler_output if return_dict else outputs[1]

     __A    :						Any						 =     self.classifier(__lowerCamelCase  )

     __A    :						int						 =     None

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

     if not return_dict:
         __A    :						Optional[Any]						 =     (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"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a_ =			logging.get_logger(__name__)

a_ =			{
    """shi-labs/nat-mini-in1k-224""": """https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json""",
    # See all Nat models at https://huggingface.co/models?filter=nat
}






class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     """nat"""

 _lowerCamelCase        =     {
     """num_attention_heads""": """num_heads""",
     """num_hidden_layers""": """num_layers""",
 }

 def __init__(    self					,						__lowerCamelCase=4					,						__lowerCamelCase=3					,						__lowerCamelCase=64					,						__lowerCamelCase=[3, 4, 6, 5]					,						__lowerCamelCase=[2, 4, 8, 16]					,						__lowerCamelCase=7					,						__lowerCamelCase=3.0					,						__lowerCamelCase=True					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.1					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=1e-5					,						__lowerCamelCase=0.0					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     super().__init__(**__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     patch_size
     __A    :						Optional[Any]						 =     num_channels
     __A    :						Tuple						 =     embed_dim
     __A    :						Dict						 =     depths
     __A    :						str						 =     len(__lowerCamelCase  )
     __A    :						Optional[Any]						 =     num_heads
     __A    :						str						 =     kernel_size
     __A    :						Any						 =     mlp_ratio
     __A    :						Optional[int]						 =     qkv_bias
     __A    :						str						 =     hidden_dropout_prob
     __A    :						Any						 =     attention_probs_dropout_prob
     __A    :						int						 =     drop_path_rate
     __A    :						int						 =     hidden_act
     __A    :						Any						 =     layer_norm_eps
     __A    :						Tuple						 =     initializer_range
     # we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
     # this indicates the channel dimension after the last stage of the model
     __A    :						int						 =     int(embed_dim * 2 ** (len(__lowerCamelCase  ) - 1)  )
     __A    :						Union[str, Any]						 =     layer_scale_init_value
     __A    :						List[str]						 =     ['''stem'''] + [F"""stage{idx}""" for idx in range(1					,						len(__lowerCamelCase  ) + 1  )]
     __A ,       __A    :						Any						 =     get_aligned_output_features_output_indices(
         out_features=__lowerCamelCase					,						out_indices=__lowerCamelCase					,						stage_names=self.stage_names  )




