


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, 1088, 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 	_a							(      nn.Module      ):

	def __init__(							self:     List[Any]   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int = 3   ,						UpperCamelCase_:     int = 1   ,						UpperCamelCase_:     int = 1   ,						UpperCamelCase_:     Optional[str] = "relu"   ,						)      ->       List[str]:





			"""simple docstring"""







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


	def 				lowerCamelCase_			(							self:     Union[str, Any]   ,						UpperCamelCase_:     List[str]							)      ->       Tuple:





			"""simple docstring"""







			lowercase__						=				self.convolution(UpperCamelCase_							)
			lowercase__						=				self.normalization(UpperCamelCase_							)
			lowercase__						=				self.activation(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Union[str, Any]   ,						UpperCamelCase_:     RegNetConfig							)      ->       Optional[Any]:





			"""simple docstring"""







			super().__init__()
			lowercase__						=				RegNetConvLayer(
			    config.num_channels   ,						config.embedding_size   ,						kernel_size=3   ,						stride=2   ,						activation=config.hidden_act							)
			lowercase__						=				config.num_channels


	def 				lowerCamelCase_			(							self:     List[str]   ,						UpperCamelCase_:     str							)      ->       Union[str, Any]:





			"""simple docstring"""







			lowercase__						=				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__						=				self.embedder(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     List[Any]   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int = 2							)      ->       List[Any]:





			"""simple docstring"""







			super().__init__()
			lowercase__						=				nn.Convad(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1   ,						stride=UpperCamelCase_   ,						bias=UpperCamelCase_							)
			lowercase__						=				nn.BatchNormad(UpperCamelCase_							)


	def 				lowerCamelCase_			(							self:     str   ,						UpperCamelCase_:     Tensor							)      ->       Tensor:





			"""simple docstring"""







			lowercase__						=				self.convolution(UpperCamelCase_							)
			lowercase__						=				self.normalization(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Tuple   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int							)      ->       List[Any]:





			"""simple docstring"""







			super().__init__()

			lowercase__						=				nn.AdaptiveAvgPoolad((1, 1)							)
			lowercase__						=				nn.Sequential(
			    nn.Convad(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1							)   ,						nn.ReLU()   ,						nn.Convad(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1							)   ,						nn.Sigmoid()   ,						)


	def 				lowerCamelCase_			(							self:     int   ,						UpperCamelCase_:     int							)      ->       Any:





			"""simple docstring"""







			lowercase__						=				self.pooler(UpperCamelCase_							)
			lowercase__						=				self.attention(UpperCamelCase_							)
			lowercase__						=				hidden_state * attention
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Tuple   ,						UpperCamelCase_:     RegNetConfig   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int = 1							)      ->       str:





			"""simple docstring"""







			super().__init__()
			lowercase__						=				in_channels != out_channels or stride != 1
			lowercase__						=				max(1   ,						out_channels // config.groups_width							)
			lowercase__						=				(
			    RegNetShortCut(UpperCamelCase_   ,						UpperCamelCase_   ,						stride=UpperCamelCase_							) if should_apply_shortcut else nn.Identity()
			)
			lowercase__						=				nn.Sequential(
			    RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1   ,						activation=config.hidden_act							)   ,						RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						stride=UpperCamelCase_   ,						groups=UpperCamelCase_   ,						activation=config.hidden_act							)   ,						RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1   ,						activation=UpperCamelCase_							)   ,						)
			lowercase__						=				ACTaFN[config.hidden_act]


	def 				lowerCamelCase_			(							self:     Any   ,						UpperCamelCase_:     Union[str, Any]							)      ->       str:





			"""simple docstring"""







			lowercase__						=				hidden_state
			lowercase__						=				self.layer(UpperCamelCase_							)
			lowercase__						=				self.shortcut(UpperCamelCase_							)
			hidden_state += residual
			lowercase__						=				self.activation(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Optional[Any]   ,						UpperCamelCase_:     RegNetConfig   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int = 1							)      ->       Tuple:





			"""simple docstring"""







			super().__init__()
			lowercase__						=				in_channels != out_channels or stride != 1
			lowercase__						=				max(1   ,						out_channels // config.groups_width							)
			lowercase__						=				(
			    RegNetShortCut(UpperCamelCase_   ,						UpperCamelCase_   ,						stride=UpperCamelCase_							) if should_apply_shortcut else nn.Identity()
			)
			lowercase__						=				nn.Sequential(
			    RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1   ,						activation=config.hidden_act							)   ,						RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						stride=UpperCamelCase_   ,						groups=UpperCamelCase_   ,						activation=config.hidden_act							)   ,						RegNetSELayer(UpperCamelCase_   ,						reduced_channels=int(round(in_channels / 4							)							)							)   ,						RegNetConvLayer(UpperCamelCase_   ,						UpperCamelCase_   ,						kernel_size=1   ,						activation=UpperCamelCase_							)   ,						)
			lowercase__						=				ACTaFN[config.hidden_act]


	def 				lowerCamelCase_			(							self:     str   ,						UpperCamelCase_:     List[str]							)      ->       int:





			"""simple docstring"""







			lowercase__						=				hidden_state
			lowercase__						=				self.layer(UpperCamelCase_							)
			lowercase__						=				self.shortcut(UpperCamelCase_							)
			hidden_state += residual
			lowercase__						=				self.activation(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Union[str, Any]   ,						UpperCamelCase_:     RegNetConfig   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int   ,						UpperCamelCase_:     int = 2   ,						UpperCamelCase_:     int = 2   ,						)      ->       Dict:





			"""simple docstring"""







			super().__init__()

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

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


	def 				lowerCamelCase_			(							self:     Union[str, Any]   ,						UpperCamelCase_:     List[Any]							)      ->       Union[str, Any]:





			"""simple docstring"""







			lowercase__						=				self.layers(UpperCamelCase_							)
			return hidden_state


class 	_a							(      nn.Module      ):

	def __init__(							self:     Dict   ,						UpperCamelCase_:     RegNetConfig							)      ->       str:





			"""simple docstring"""







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


	def 				lowerCamelCase_			(							self:     Tuple   ,						UpperCamelCase_:     Tensor   ,						UpperCamelCase_:     bool = False   ,						UpperCamelCase_:     bool = True							)      ->       BaseModelOutputWithNoAttention:





			"""simple docstring"""







			lowercase__						=				() if output_hidden_states else None

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

					lowercase__						=				stage_module(UpperCamelCase_							)

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


class 	_a							(      UpperCamelCase__      ):
	_lowercase     :      Tuple  				=			RegNetConfig
	_lowercase     :      str  				=			'''regnet'''
	_lowercase     :      Tuple  				=			'''pixel_values'''
	_lowercase     :      Dict  				=			True

	def 				lowerCamelCase_			(							self:     int   ,						UpperCamelCase_:     Optional[Any]							)      ->       str:





			"""simple docstring"""







			if isinstance(UpperCamelCase_   ,						nn.Convad							):
					nn.init.kaiming_normal_(module.weight   ,						mode='''fan_out'''   ,						nonlinearity='''relu'''							)
			elif isinstance(UpperCamelCase_   ,						(nn.BatchNormad, nn.GroupNorm)							):
					nn.init.constant_(module.weight   ,						1							)
					nn.init.constant_(module.bias   ,						0							)


	def 				lowerCamelCase_			(							self:     Any   ,						UpperCamelCase_:     List[Any]   ,						UpperCamelCase_:     Tuple=False							)      ->       Optional[int]:





			"""simple docstring"""







			if isinstance(UpperCamelCase_   ,						UpperCamelCase_							):
					lowercase__						=				value
lowerCAmelCase	=							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'

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

	def __init__(							self:     List[str]   ,						UpperCamelCase_:     List[Any]							)      ->       Any:





			"""simple docstring"""







			super().__init__(UpperCamelCase_							)
			lowercase__						=				config
			lowercase__						=				RegNetEmbeddings(UpperCamelCase_							)
			lowercase__						=				RegNetEncoder(UpperCamelCase_							)
			lowercase__						=				nn.AdaptiveAvgPoolad((1, 1)							)
			# Initialize weights and apply final processing
			self.post_init()


	@add_start_docstrings_to_model_forward(UpperCamelCase_							)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC   ,						output_type=UpperCamelCase_   ,						config_class=_CONFIG_FOR_DOC   ,						modality='''vision'''   ,						expected_output=_EXPECTED_OUTPUT_SHAPE   ,						)
	def 				lowerCamelCase_			(							self:     Optional[Any]   ,						UpperCamelCase_:     Tensor   ,						UpperCamelCase_:     Optional[bool] = None   ,						UpperCamelCase_:     Optional[bool] = None							)      ->       BaseModelOutputWithPoolingAndNoAttention:





			"""simple docstring"""







			lowercase__						=				(
			    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
			)
			lowercase__						=				return_dict if return_dict is not None else self.config.use_return_dict

			lowercase__						=				self.embedder(UpperCamelCase_							)

			lowercase__						=				self.encoder(
			    UpperCamelCase_   ,						output_hidden_states=UpperCamelCase_   ,						return_dict=UpperCamelCase_							)

			lowercase__						=				encoder_outputs[0]

			lowercase__						=				self.pooler(UpperCamelCase_							)

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

			return BaseModelOutputWithPoolingAndNoAttention(
			    last_hidden_state=UpperCamelCase_   ,						pooler_output=UpperCamelCase_   ,						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.
    '''					,     UpperCamelCase__					,     )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class 	_a							(      UpperCamelCase__      ):

	def __init__(							self:     Union[str, Any]   ,						UpperCamelCase_:     List[str]							)      ->       Optional[Any]:





			"""simple docstring"""







			super().__init__(UpperCamelCase_							)
			lowercase__						=				config.num_labels
			lowercase__						=				RegNetModel(UpperCamelCase_							)
			# classification head
			lowercase__						=				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(UpperCamelCase_							)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,						output_type=UpperCamelCase_   ,						config_class=_CONFIG_FOR_DOC   ,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,						)
	def 				lowerCamelCase_			(							self:     Tuple   ,						UpperCamelCase_:     Optional[torch.FloatTensor] = None   ,						UpperCamelCase_:     Optional[torch.LongTensor] = None   ,						UpperCamelCase_:     Optional[bool] = None   ,						UpperCamelCase_:     Optional[bool] = None   ,						)      ->       ImageClassifierOutputWithNoAttention:





			"""simple docstring"""







			lowercase__						=				return_dict if return_dict is not None else self.config.use_return_dict

			lowercase__						=				self.regnet(UpperCamelCase_   ,						output_hidden_states=UpperCamelCase_   ,						return_dict=UpperCamelCase_							)

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

			lowercase__						=				self.classifier(UpperCamelCase_							)

			lowercase__						=				None

			if labels is not None:
					if self.config.problem_type is None:
							if self.num_labels == 1:
									lowercase__						=				'''regression'''
							elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
									lowercase__						=				'''single_label_classification'''
							else:
									lowercase__						=				'''multi_label_classification'''
					if self.config.problem_type == "regression":
							lowercase__						=				MSELoss()
							if self.num_labels == 1:
									lowercase__						=				loss_fct(logits.squeeze()   ,						labels.squeeze()							)
							else:
									lowercase__						=				loss_fct(UpperCamelCase_   ,						UpperCamelCase_							)
					elif self.config.problem_type == "single_label_classification":
							lowercase__						=				CrossEntropyLoss()
							lowercase__						=				loss_fct(logits.view(-1   ,						self.num_labels							)   ,						labels.view(-1							)							)
					elif self.config.problem_type == "multi_label_classification":
							lowercase__						=				BCEWithLogitsLoss()
							lowercase__						=				loss_fct(UpperCamelCase_   ,						UpperCamelCase_							)

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

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






# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


lowerCAmelCase	=							'Create a default config file for Accelerate with only a few flags set.'
def       _a			(				SCREAMING_SNAKE_CASE="no"						,    SCREAMING_SNAKE_CASE = default_json_config_file						,    SCREAMING_SNAKE_CASE = False					):


		"""simple docstring"""




		lowercase__						=				Path(SCREAMING_SNAKE_CASE					)
		path.parent.mkdir(parents=SCREAMING_SNAKE_CASE						,    exist_ok=SCREAMING_SNAKE_CASE					)
		if path.exists():
				print(
				    f'Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'					)
				return False
		lowercase__						=				mixed_precision.lower()
		if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
				raise ValueError(
				    f'`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'					)
		lowercase__						=				{
		    '''compute_environment''': '''LOCAL_MACHINE''',
		    '''mixed_precision''': mixed_precision,
		}
		if torch.cuda.is_available():
				lowercase__						=				torch.cuda.device_count()
				lowercase__						=				num_gpus
				lowercase__						=				False
				if num_gpus > 1:
						lowercase__						=				'''MULTI_GPU'''
				else:
						lowercase__						=				'''NO'''
		elif is_xpu_available() and use_xpu:
				lowercase__						=				torch.xpu.device_count()
				lowercase__						=				num_xpus
				lowercase__						=				False
				if num_xpus > 1:
						lowercase__						=				'''MULTI_XPU'''
				else:
						lowercase__						=				'''NO'''
		elif is_npu_available():
				lowercase__						=				torch.npu.device_count()
				lowercase__						=				num_npus
				lowercase__						=				False
				if num_npus > 1:
						lowercase__						=				'''MULTI_NPU'''
				else:
						lowercase__						=				'''NO'''
		else:
				lowercase__						=				0
				lowercase__						=				True
				lowercase__						=				1
				lowercase__						=				'''NO'''
		lowercase__						=				ClusterConfig(**SCREAMING_SNAKE_CASE					)
		config.to_json_file(SCREAMING_SNAKE_CASE					)
		return path
def       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				parser.add_parser('''default'''						,    parents=SCREAMING_SNAKE_CASE						,    help=SCREAMING_SNAKE_CASE						,    formatter_class=SCREAMING_SNAKE_CASE					)
		parser.add_argument(
		    '''--config_file'''						,    default=SCREAMING_SNAKE_CASE						,    help=(
		        '''The path to use to store the config file. Will default to a file named default_config.yaml in the cache '''
		        '''location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '''
		        '''such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '''
		        '''with \'huggingface\'.'''
		    )						,    dest='''save_location'''						,    )

		parser.add_argument(
		    '''--mixed_precision'''						,    choices=['''no''', '''fp16''', '''bf16''']						,    type=SCREAMING_SNAKE_CASE						,    help='''Whether or not to use mixed precision training. '''
		    '''Choose between FP16 and BF16 (bfloat16) training. '''
		    '''BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.'''						,    default='''no'''						,    )
		parser.set_defaults(func=SCREAMING_SNAKE_CASE					)
		return parser
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				write_basic_config(args.mixed_precision						,    args.save_location					)
		if config_file:
				print(f'accelerate configuration saved at {config_file}'					)


