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


__lowerCAmelCase         =    logging.get_logger(__name__)

# General docstring
__lowerCAmelCase         =    """RegNetConfig"""

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

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

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







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

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


     '''simple docstring'''






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


  def   __lowercase		(		self      :     List[str]							,_a      :     int       ):


     '''simple docstring'''






     _a				:	Dict			=   self.convolution(_a       )
     _a				:	int			=   self.normalization(_a       )
     _a				:	int			=   self.activation(_a       )
     return hidden_state







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     List[str]							,_a      :     RegNetConfig       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     List[Any]							,_a      :     int       ):


     '''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				:	Any			=   self.embedder(_a       )
     return hidden_state







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     List[Any]							,_a      :     int							,_a      :     int							,_a      :     int = 2       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     str							,_a      :     Tensor       ):


     '''simple docstring'''






     _a				:	int			=   self.convolution(_a       )
     _a				:	Optional[Any]			=   self.normalization(_a       )
     return hidden_state







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

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


     '''simple docstring'''






     super().__init__()

     _a				:	str			=   nn.AdaptiveAvgPoolad((1, 1)       )
     _a				:	Union[str, Any]			=   nn.Sequential(
         nn.Convad(_a							,_a							,kernel_size=1       )							,nn.ReLU()							,nn.Convad(_a							,_a							,kernel_size=1       )							,nn.Sigmoid()							,)


  def   __lowercase		(		self      :     Union[str, Any]							,_a      :     str       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.pooler(_a       )
     _a				:	Tuple			=   self.attention(_a       )
     _a				:	Optional[Any]			=   hidden_state * attention
     return hidden_state







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     Any							,_a      :     RegNetConfig							,_a      :     int							,_a      :     int							,_a      :     int = 1       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Any							,_a      :     int       ):


     '''simple docstring'''






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







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     Dict							,_a      :     RegNetConfig							,_a      :     int							,_a      :     int							,_a      :     int = 1       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Dict							,_a      :     Union[str, Any]       ):


     '''simple docstring'''






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







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     List[str]							,_a      :     RegNetConfig							,_a      :     int							,_a      :     int							,_a      :     int = 2							,_a      :     int = 2							,):


     '''simple docstring'''






     super().__init__()

     _a				:	List[str]			=   RegNetXLayer if config.layer_type == 'x' else RegNetYLayer

     _a				:	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      :     Any							,_a      :     Dict       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   self.layers(_a       )
     return hidden_state







class UpperCAmelCase__	( nn.Module	):





  """simple docstring"""

  def __init__(		self      :     Any							,_a      :     RegNetConfig       ):


     '''simple docstring'''






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


     '''simple docstring'''






     _a				:	List[str]			=   () if output_hidden_states else None

     for stage_module in self.stages:
        if output_hidden_states:
           _a				:	Union[str, Any]			=   hidden_states + (hidden_state,)

        _a				:	int			=   stage_module(_a       )

     if output_hidden_states:
        _a				:	Any			=   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 UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  __UpperCAmelCase				:				int   =					RegNetConfig
  __UpperCAmelCase				:				Any   =					'''regnet'''
  __UpperCAmelCase				:				Optional[Any]   =					'''pixel_values'''
  __UpperCAmelCase				:				List[str]   =					True


  def   __lowercase		(		self      :     str							,_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      :     Union[str, Any]							,_a      :     int							,_a      :     List[str]=False       ):


     '''simple docstring'''






     if isinstance(_a							,_a       ):
        _a				:	Optional[int]			=   value


__lowerCAmelCase         =    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.
"""

__lowerCAmelCase         =    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.'''    ,				lowercase__    ,				)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     int							,_a      :     int       ):


     '''simple docstring'''






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


     '''simple docstring'''






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

     _a				:	Optional[Any]			=   self.embedder(_a       )

     _a				:	Dict			=   self.encoder(
         _a							,output_hidden_states=_a							,return_dict=_a       )

     _a				:	Tuple			=   encoder_outputs[0]

     _a				:	Dict			=   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.
    '''    ,				lowercase__    ,				)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     str							,_a      :     Optional[Any]       ):


     '''simple docstring'''






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


     '''simple docstring'''






     _a				:	Any			=   return_dict if return_dict is not None else self.config.use_return_dict

     _a				:	int			=   self.regnet(_a							,output_hidden_states=_a							,return_dict=_a       )

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

     _a				:	Optional[Any]			=   self.classifier(_a       )

     _a				:	Dict			=   None

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

     if not return_dict:
        _a				:	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'''







import warnings

from ...utils import logging
from .image_processing_videomae import VideoMAEImageProcessor


__lowerCAmelCase         =    logging.get_logger(__name__)







class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     Tuple							,*_a      :     List[str]							,**_a      :     Any       ):


     '''simple docstring'''






     warnings.warn(
         'The class VideoMAEFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
         ' Please use VideoMAEImageProcessor instead.'							,_a							,)
     super().__init__(*_a							,**_a       )



