





'''simple docstring'''



from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


_SCREAMING_SNAKE_CASE					=    logging.get_logger(__name__)


# General docstring
_SCREAMING_SNAKE_CASE					=    '''MobileNetV1Config'''

# Base docstring
_SCREAMING_SNAKE_CASE					=    '''google/mobilenet_v1_1.0_224'''
_SCREAMING_SNAKE_CASE					=    [1, 1_0_2_4, 7, 7]

# Image classification docstring
_SCREAMING_SNAKE_CASE					=    '''google/mobilenet_v1_1.0_224'''
_SCREAMING_SNAKE_CASE					=    '''tabby, tabby cat'''


_SCREAMING_SNAKE_CASE					=    [
    '''google/mobilenet_v1_1.0_224''',
    '''google/mobilenet_v1_0.75_192''',
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]




def    _lowerCAmelCase						(				lowerCamelCase_			:  str		,       lowerCamelCase_			:  Union[str, Any]		,       lowerCamelCase_			:  List[Any]=None				):
			__lowercase										=							{}

			if isinstance(lowerCamelCase_		,       lowerCamelCase_				):
						__lowercase										=							model.mobilenet_va
			else:
						__lowercase										=							model

			__lowercase										=							'''MobilenetV1/Conv2d_0/'''
			__lowercase										=							backbone.conv_stem.convolution.weight
			__lowercase										=							backbone.conv_stem.normalization.bias
			__lowercase										=							backbone.conv_stem.normalization.weight
			__lowercase										=							backbone.conv_stem.normalization.running_mean
			__lowercase										=							backbone.conv_stem.normalization.running_var

			for i in range(1_3				):
						__lowercase										=							i + 1
						__lowercase										=							i * 2

						__lowercase										=							backbone.layer[pt_index]
						__lowercase										=							f"MobilenetV1/Conv2d_{tf_index}_depthwise/"
						__lowercase										=							pointer.convolution.weight
						__lowercase										=							pointer.normalization.bias
						__lowercase										=							pointer.normalization.weight
						__lowercase										=							pointer.normalization.running_mean
						__lowercase										=							pointer.normalization.running_var

						__lowercase										=							backbone.layer[pt_index + 1]
						__lowercase										=							f"MobilenetV1/Conv2d_{tf_index}_pointwise/"
						__lowercase										=							pointer.convolution.weight
						__lowercase										=							pointer.normalization.bias
						__lowercase										=							pointer.normalization.weight
						__lowercase										=							pointer.normalization.running_mean
						__lowercase										=							pointer.normalization.running_var

			if isinstance(lowerCamelCase_		,       lowerCamelCase_				):
						__lowercase										=							'''MobilenetV1/Logits/Conv2d_1c_1x1/'''
						__lowercase										=							model.classifier.weight
						__lowercase										=							model.classifier.bias

			return tf_to_pt_map




def    _lowerCAmelCase						(				lowerCamelCase_			:  Tuple		,       lowerCamelCase_			:  Optional[int]		,       lowerCamelCase_			:  Tuple				):
			try:
						import numpy as np
						import tensorflow as tf
			except ImportError:
						logger.error(
						    '''Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see '''
						    '''https://www.tensorflow.org/install/ for installation instructions.'''				)
						raise

			# Load weights from TF model
			__lowercase										=							tf.train.list_variables(lowerCamelCase_				)
			__lowercase										=							{}
			for name, shape in init_vars:
						logger.info(f"Loading TF weight {name} with shape {shape}"				)
						__lowercase										=							tf.train.load_variable(lowerCamelCase_		,       lowerCamelCase_				)
						__lowercase										=							array

			# Build TF to PyTorch weights loading map
			__lowercase										=							_build_tf_to_pytorch_map(lowerCamelCase_		,       lowerCamelCase_		,       lowerCamelCase_				)

			for name, pointer in tf_to_pt_map.items():
						logger.info(f"Importing {name}"				)
						if name not in tf_weights:
									logger.info(f"{name} not in tf pre-trained weights, skipping"				)
									continue

						__lowercase										=							tf_weights[name]

						if "depthwise_weights" in name:
									logger.info('''Transposing depthwise'''				)
									__lowercase										=							np.transpose(lowerCamelCase_		,       (2, 3, 0, 1)				)
						elif "weights" in name:
									logger.info('''Transposing'''				)
									if len(pointer.shape				) == 2:  # copying into linear layer
												__lowercase										=							array.squeeze().transpose()
									else:
												__lowercase										=							np.transpose(lowerCamelCase_		,       (3, 2, 0, 1)				)

						if pointer.shape != array.shape:
									raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"				)

						logger.info(f"Initialize PyTorch weight {name} {array.shape}"				)
						__lowercase										=							torch.from_numpy(lowerCamelCase_				)

						tf_weights.pop(lowerCamelCase_		,       lowerCamelCase_				)
						tf_weights.pop(name + '''/RMSProp'''		,       lowerCamelCase_				)
						tf_weights.pop(name + '''/RMSProp_1'''		,       lowerCamelCase_				)
						tf_weights.pop(name + '''/ExponentialMovingAverage'''		,       lowerCamelCase_				)

			logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys()				)}"				)
			return model




def    _lowerCAmelCase						(				lowerCamelCase_			:  torch.Tensor		,       lowerCamelCase_			:  nn.Convad				):
			__lowercase						,		__lowercase										=							features.shape[-2:]
			__lowercase						,		__lowercase										=							conv_layer.stride
			__lowercase						,		__lowercase										=							conv_layer.kernel_size

			if in_height % stride_height == 0:
						__lowercase										=							max(kernel_height - stride_height		,       0				)
			else:
						__lowercase										=							max(kernel_height - (in_height % stride_height)		,       0				)

			if in_width % stride_width == 0:
						__lowercase										=							max(kernel_width - stride_width		,       0				)
			else:
						__lowercase										=							max(kernel_width - (in_width % stride_width)		,       0				)

			__lowercase										=							pad_along_width // 2
			__lowercase										=							pad_along_width - pad_left
			__lowercase										=							pad_along_height // 2
			__lowercase										=							pad_along_height - pad_top

			__lowercase										=							(pad_left, pad_right, pad_top, pad_bottom)
			return nn.functional.pad(lowerCamelCase_		,       lowerCamelCase_		,       '''constant'''		,       0.0				)






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 1     ,_lowerCamelCase = 1     ,_lowerCamelCase = False     ,_lowerCamelCase = True     ,_lowerCamelCase = True     ,)							->							None:







				'''simple docstring'''





				super().__init__()
				__lowercase										=							config

				if in_channels % groups != 0:
							raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups." )
				if out_channels % groups != 0:
							raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups." )

				__lowercase										=							0 if config.tf_padding else int((kernel_size - 1) / 2 )

				__lowercase										=							nn.Convad(
				    in_channels=_lowerCamelCase     ,out_channels=_lowerCamelCase     ,kernel_size=_lowerCamelCase     ,stride=_lowerCamelCase     ,padding=_lowerCamelCase     ,groups=_lowerCamelCase     ,bias=_lowerCamelCase     ,padding_mode='''zeros'''     ,)

				if use_normalization:
							__lowercase										=							nn.BatchNormad(
							    num_features=_lowerCamelCase     ,eps=config.layer_norm_eps     ,momentum=0.9_9_9_7     ,affine=_lowerCamelCase     ,track_running_stats=_lowerCamelCase     ,)
				else:
							__lowercase										=							None

				if use_activation:
							if isinstance(_lowerCamelCase     ,_lowerCamelCase ):
										__lowercase										=							ACTaFN[use_activation]
							elif isinstance(config.hidden_act     ,_lowerCamelCase ):
										__lowercase										=							ACTaFN[config.hidden_act]
							else:
										__lowercase										=							config.hidden_act
				else:
							__lowercase										=							None


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							torch.Tensor:







				'''simple docstring'''





				if self.config.tf_padding:
							__lowercase										=							apply_tf_padding(_lowerCamelCase     ,self.convolution )
				__lowercase										=							self.convolution(_lowerCamelCase )
				if self.normalization is not None:
							__lowercase										=							self.normalization(_lowerCamelCase )
				if self.activation is not None:
							__lowercase										=							self.activation(_lowerCamelCase )
				return features






class 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	a    :				Optional[int]											=       MobileNetVaConfig
	a    :				Union[str, Any]											=       load_tf_weights_in_mobilenet_va
	a    :				Optional[Any]											=       "mobilenet_v1"
	a    :				int											=       "pixel_values"
	a    :				Union[str, Any]											=       False


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							None:







				'''simple docstring'''





				if isinstance(_lowerCamelCase     ,(nn.Linear, nn.Convad) ):
							module.weight.data.normal_(mean=0.0     ,std=self.config.initializer_range )
							if module.bias is not None:
										module.bias.data.zero_()
				elif isinstance(_lowerCamelCase     ,nn.BatchNormad ):
							module.bias.data.zero_()
							module.weight.data.fill_(1.0 )





_SCREAMING_SNAKE_CASE					=    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 ([`MobileNetV1Config`]): 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.
'''

_SCREAMING_SNAKE_CASE					=    R'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`MobileNetV1ImageProcessor.__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 [`~utils.ModelOutput`] instead of a plain tuple.
'''






@add_start_docstrings(
    "The bare MobileNetV1 model outputting raw hidden-states without any specific head on top." ,      lowerCAmelCase__ ,      )
class 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase = True )							->							str:







				'''simple docstring'''





				super().__init__(_lowerCamelCase )
				__lowercase										=							config

				__lowercase										=							32
				__lowercase										=							max(int(depth * config.depth_multiplier )     ,config.min_depth )

				__lowercase										=							MobileNetVaConvLayer(
				    _lowerCamelCase     ,in_channels=config.num_channels     ,out_channels=_lowerCamelCase     ,kernel_size=3     ,stride=2     ,)

				__lowercase										=							[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

				__lowercase										=							nn.ModuleList()
				for i in range(13 ):
							__lowercase										=							out_channels

							if strides[i] == 2 or i == 0:
										depth *= 2
										__lowercase										=							max(int(depth * config.depth_multiplier )     ,config.min_depth )

							self.layer.append(
							    MobileNetVaConvLayer(
							        _lowerCamelCase     ,in_channels=_lowerCamelCase     ,out_channels=_lowerCamelCase     ,kernel_size=3     ,stride=strides[i]     ,groups=_lowerCamelCase     ,) )

							self.layer.append(
							    MobileNetVaConvLayer(
							        _lowerCamelCase     ,in_channels=_lowerCamelCase     ,out_channels=_lowerCamelCase     ,kernel_size=1     ,) )

				__lowercase										=							nn.AdaptiveAvgPoolad((1, 1) ) if add_pooling_layer else None

				# Initialize weights and apply final processing
				self.post_init()


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							Tuple:







				'''simple docstring'''





				raise NotImplementedError


	@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 = None     ,_lowerCamelCase = None     ,_lowerCamelCase = None     ,)							->							Union[tuple, 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

				if pixel_values is None:
							raise ValueError('''You have to specify pixel_values''' )

				__lowercase										=							self.conv_stem(_lowerCamelCase )

				__lowercase										=							() if output_hidden_states else None

				for i, layer_module in enumerate(self.layer ):
							__lowercase										=							layer_module(_lowerCamelCase )

							if output_hidden_states:
										__lowercase										=							all_hidden_states + (hidden_states,)

				__lowercase										=							hidden_states

				if self.pooler is not None:
							__lowercase										=							torch.flatten(self.pooler(_lowerCamelCase )     ,start_dim=1 )
				else:
							__lowercase										=							None

				if not return_dict:
							return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None )

				return BaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=_lowerCamelCase     ,pooler_output=_lowerCamelCase     ,hidden_states=_lowerCamelCase     ,)






@add_start_docstrings(
    "\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    " ,      lowerCAmelCase__ ,      )
class 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase )							->							None:







				'''simple docstring'''





				super().__init__(_lowerCamelCase )

				__lowercase										=							config.num_labels
				__lowercase										=							MobileNetVaModel(_lowerCamelCase )

				__lowercase										=							self.mobilenet_va.layer[-1].convolution.out_channels

				# Classifier head
				__lowercase										=							nn.Dropout(config.classifier_dropout_prob     ,inplace=_lowerCamelCase )
				__lowercase										=							nn.Linear(_lowerCamelCase     ,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     ,)							->							Union[tuple, ImageClassifierOutputWithNoAttention]:







				'''simple docstring'''





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

				__lowercase										=							self.mobilenet_va(_lowerCamelCase     ,output_hidden_states=_lowerCamelCase     ,return_dict=_lowerCamelCase )

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

				__lowercase										=							self.classifier(self.dropout(_lowerCamelCase ) )

				__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(_lowerCamelCase     ,_lowerCamelCase )
							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(_lowerCamelCase     ,_lowerCamelCase )

				if not return_dict:
							__lowercase										=							(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 typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_SCREAMING_SNAKE_CASE					=    {'''configuration_van''': ['''VAN_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VanConfig''']}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE					=    [
	    '''VAN_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''VanForImageClassification''',
	    '''VanModel''',
	    '''VanPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_van import (
		    VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
		    VanForImageClassification,
		    VanModel,
		    VanPreTrainedModel,
		)

else:
	import sys

	_SCREAMING_SNAKE_CASE					=    _LazyModule(__name__, globals()['''__file__'''], _import_structure)





