






'''simple docstring'''

from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


SCREAMING_SNAKE_CASE__     =       [
    'openmmlab/upernet-convnext-tiny',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
SCREAMING_SNAKE_CASE__     =       'UperNetConfig'






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = 0      ,		_SCREAMING_SNAKE_CASE = False      ,		_SCREAMING_SNAKE_CASE = 1      ,		)  ->				None:



        """simple docstring"""






        super().__init__()
        UpperCamelCase      = nn.Convad(
            in_channels=_SCREAMING_SNAKE_CASE      ,		out_channels=_SCREAMING_SNAKE_CASE      ,		kernel_size=_SCREAMING_SNAKE_CASE      ,		padding=_SCREAMING_SNAKE_CASE      ,		bias=_SCREAMING_SNAKE_CASE      ,		dilation=_SCREAMING_SNAKE_CASE      ,		)
        UpperCamelCase      = nn.BatchNormad(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = nn.ReLU()


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = self.conv(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.batch_norm(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.activation(_SCREAMING_SNAKE_CASE				)

        return output






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				None:



        """simple docstring"""






        super().__init__()
        UpperCamelCase      = [
            nn.AdaptiveAvgPoolad(_SCREAMING_SNAKE_CASE				),
            UperNetConvModule(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		kernel_size=1				),
        ]
        for i, layer in enumerate(self.layers				):
              self.add_module(str(_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = input
        for layer in self.layers:
              UpperCamelCase      = layer(_SCREAMING_SNAKE_CASE				)
        return hidden_state






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				None:



        """simple docstring"""






        super().__init__()
        UpperCamelCase      = pool_scales
        UpperCamelCase      = align_corners
        UpperCamelCase      = in_channels
        UpperCamelCase      = channels
        UpperCamelCase      = []
        for i, pool_scale in enumerate(_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = UperNetPyramidPoolingBlock(pool_scale=_SCREAMING_SNAKE_CASE      ,		in_channels=_SCREAMING_SNAKE_CASE      ,		channels=_SCREAMING_SNAKE_CASE				)
              self.blocks.append(_SCREAMING_SNAKE_CASE				)
              self.add_module(str(_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				List[torch.Tensor]:



        """simple docstring"""






        UpperCamelCase      = []
        for ppm in self.blocks:
              UpperCamelCase      = ppm(_SCREAMING_SNAKE_CASE				)
              UpperCamelCase      = nn.functional.interpolate(
                  _SCREAMING_SNAKE_CASE      ,		size=x.size()[2:]      ,		mode="""bilinear"""      ,		align_corners=self.align_corners				)
              ppm_outs.append(_SCREAMING_SNAKE_CASE				)
        return ppm_outs






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				Optional[int]:



        """simple docstring"""






        super().__init__()

        UpperCamelCase      = config
        UpperCamelCase      = config.pool_scales  # e.g. (1, 2, 3, 6)
        UpperCamelCase      = in_channels
        UpperCamelCase      = config.hidden_size
        UpperCamelCase      = False
        UpperCamelCase      = nn.Convad(self.channels      ,		config.num_labels      ,		kernel_size=1				)

        # PSP Module
        UpperCamelCase      = UperNetPyramidPoolingModule(
            self.pool_scales      ,		self.in_channels[-1]      ,		self.channels      ,		align_corners=self.align_corners      ,		)
        UpperCamelCase      = UperNetConvModule(
            self.in_channels[-1] + len(self.pool_scales				) * self.channels      ,		self.channels      ,		kernel_size=3      ,		padding=1      ,		)
        # FPN Module
        UpperCamelCase      = nn.ModuleList()
        UpperCamelCase      = nn.ModuleList()
        for in_channels in self.in_channels[:-1]:  # skip the top layer
              UpperCamelCase      = UperNetConvModule(_SCREAMING_SNAKE_CASE      ,		self.channels      ,		kernel_size=1				)
              UpperCamelCase      = UperNetConvModule(self.channels      ,		self.channels      ,		kernel_size=3      ,		padding=1				)
              self.lateral_convs.append(_SCREAMING_SNAKE_CASE				)
              self.fpn_convs.append(_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = UperNetConvModule(
            len(self.in_channels				) * self.channels      ,		self.channels      ,		kernel_size=3      ,		padding=1      ,		)

  def   A__						(				self				)  ->				Tuple:



        """simple docstring"""






        self.apply(self._init_weights				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Dict:



        """simple docstring"""






        if isinstance(_SCREAMING_SNAKE_CASE      ,		nn.Convad				):
              module.weight.data.normal_(mean=0.0      ,		std=self.config.initializer_range				)
              if module.bias is not None:
                    module.bias.data.zero_()

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				int:



        """simple docstring"""






        UpperCamelCase      = inputs[-1]
        UpperCamelCase      = [x]
        psp_outs.extend(self.psp_modules(_SCREAMING_SNAKE_CASE				)				)
        UpperCamelCase      = torch.cat(_SCREAMING_SNAKE_CASE      ,		dim=1				)
        UpperCamelCase      = self.bottleneck(_SCREAMING_SNAKE_CASE				)

        return output


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = [lateral_conv(encoder_hidden_states[i]				) for i, lateral_conv in enumerate(self.lateral_convs				)]

        laterals.append(self.psp_forward(_SCREAMING_SNAKE_CASE				)				)

        # build top-down path
        UpperCamelCase      = len(_SCREAMING_SNAKE_CASE				)
        for i in range(used_backbone_levels - 1      ,		0      ,		-1				):
              UpperCamelCase      = laterals[i - 1].shape[2:]
              UpperCamelCase      = laterals[i - 1] + nn.functional.interpolate(
                  laterals[i]      ,		size=_SCREAMING_SNAKE_CASE      ,		mode="""bilinear"""      ,		align_corners=self.align_corners				)

        # build outputs
        UpperCamelCase      = [self.fpn_convs[i](laterals[i]				) for i in range(used_backbone_levels - 1				)]
        # append psp feature
        fpn_outs.append(laterals[-1]				)

        for i in range(used_backbone_levels - 1      ,		0      ,		-1				):
              UpperCamelCase      = nn.functional.interpolate(
                  fpn_outs[i]      ,		size=fpn_outs[0].shape[2:]      ,		mode="""bilinear"""      ,		align_corners=self.align_corners				)
        UpperCamelCase      = torch.cat(_SCREAMING_SNAKE_CASE      ,		dim=1				)
        UpperCamelCase      = self.fpn_bottleneck(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = self.classifier(_SCREAMING_SNAKE_CASE				)

        return output






class     a_						(      nn.Module    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = 2      ,		_SCREAMING_SNAKE_CASE = 3      ,		_SCREAMING_SNAKE_CASE = 1				)  ->				None:



        """simple docstring"""






        super().__init__()

        UpperCamelCase      = config
        UpperCamelCase      = config.auxiliary_in_channels
        UpperCamelCase      = config.auxiliary_channels
        UpperCamelCase      = config.auxiliary_num_convs
        UpperCamelCase      = config.auxiliary_concat_input
        UpperCamelCase      = in_index

        UpperCamelCase      = (kernel_size // 2) * dilation
        UpperCamelCase      = []
        convs.append(
            UperNetConvModule(
                self.in_channels      ,		self.channels      ,		kernel_size=_SCREAMING_SNAKE_CASE      ,		padding=_SCREAMING_SNAKE_CASE      ,		dilation=_SCREAMING_SNAKE_CASE				)				)
        for i in range(self.num_convs - 1				):
              convs.append(
                  UperNetConvModule(
                      self.channels      ,		self.channels      ,		kernel_size=_SCREAMING_SNAKE_CASE      ,		padding=_SCREAMING_SNAKE_CASE      ,		dilation=_SCREAMING_SNAKE_CASE				)				)
        if self.num_convs == 0:
              UpperCamelCase      = nn.Identity()
        else:
              UpperCamelCase      = nn.Sequential(*_SCREAMING_SNAKE_CASE				)
        if self.concat_input:
              UpperCamelCase      = UperNetConvModule(
                  self.in_channels + self.channels      ,		self.channels      ,		kernel_size=_SCREAMING_SNAKE_CASE      ,		padding=kernel_size // 2				)

        UpperCamelCase      = nn.Convad(self.channels      ,		config.num_labels      ,		kernel_size=1				)

  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






        self.apply(self._init_weights				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Optional[int]:



        """simple docstring"""






        if isinstance(_SCREAMING_SNAKE_CASE      ,		nn.Convad				):
              module.weight.data.normal_(mean=0.0      ,		std=self.config.initializer_range				)
              if module.bias is not None:
                    module.bias.data.zero_()


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				torch.Tensor:



        """simple docstring"""






        UpperCamelCase      = encoder_hidden_states[self.in_index]
        UpperCamelCase      = self.convs(_SCREAMING_SNAKE_CASE				)
        if self.concat_input:
              UpperCamelCase      = self.conv_cat(torch.cat([hidden_states, output]      ,		dim=1				)				)
        UpperCamelCase      = self.classifier(_SCREAMING_SNAKE_CASE				)
        return output






class     a_						(      lowerCamelCase    ):
  lowercase							=				UperNetConfig
  lowercase							=				"""pixel_values"""
  lowercase							=				True

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				str:



        """simple docstring"""






        if isinstance(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				):
              module.backbone.init_weights()
              module.decode_head.init_weights()
              module.auxiliary_head.init_weights()

  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






        self.backbone.init_weights()
        self.decode_head.init_weights()
        self.auxiliary_head.init_weights()


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE=False				)  ->				Union[str, Any]:



        """simple docstring"""






        if isinstance(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = value


SCREAMING_SNAKE_CASE__     =       R'\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): 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'

SCREAMING_SNAKE_CASE__     =       R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'






@add_start_docstrings(
    """UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."""     ,	lowerCamelCase     ,	)
class     a_						(      lowerCamelCase    ):

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



        """simple docstring"""






        super().__init__(_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = AutoBackbone.from_config(config.backbone_config				)

        # Semantic segmentation head(s)
        UpperCamelCase      = UperNetHead(_SCREAMING_SNAKE_CASE      ,		in_channels=self.backbone.channels				)
        UpperCamelCase      = UperNetFCNHead(_SCREAMING_SNAKE_CASE				) if config.use_auxiliary_head else None

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


  @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""				)				)
  @replace_return_docstrings(output_type=_SCREAMING_SNAKE_CASE      ,		config_class=_CONFIG_FOR_DOC				)
  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		)  ->				Union[tuple, SemanticSegmenterOutput]:



        """simple docstring"""






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

        UpperCamelCase      = self.backbone.forward_with_filtered_kwargs(
            _SCREAMING_SNAKE_CASE      ,		output_hidden_states=_SCREAMING_SNAKE_CASE      ,		output_attentions=_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = outputs.feature_maps

        UpperCamelCase      = self.decode_head(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = nn.functional.interpolate(_SCREAMING_SNAKE_CASE      ,		size=pixel_values.shape[2:]      ,		mode="""bilinear"""      ,		align_corners=_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = None
        if self.auxiliary_head is not None:
              UpperCamelCase      = self.auxiliary_head(_SCREAMING_SNAKE_CASE				)
              UpperCamelCase      = nn.functional.interpolate(
                  _SCREAMING_SNAKE_CASE      ,		size=pixel_values.shape[2:]      ,		mode="""bilinear"""      ,		align_corners=_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = None
        if labels is not None:
              if self.config.num_labels == 1:
                    raise ValueError("""The number of labels should be greater than one"""				)
              else:
                    # compute weighted loss
                    UpperCamelCase      = CrossEntropyLoss(ignore_index=self.config.loss_ignore_index				)
                    UpperCamelCase      = loss_fct(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
                    UpperCamelCase      = loss_fct(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
                    UpperCamelCase      = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

        return SemanticSegmenterOutput(
            loss=_SCREAMING_SNAKE_CASE      ,		logits=_SCREAMING_SNAKE_CASE      ,		hidden_states=outputs.hidden_states      ,		attentions=outputs.attentions      ,		)








'''simple docstring'''

import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import List

import timm
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from torch import Tensor

from transformers import AutoImageProcessor, ResNetConfig, ResNetForImageClassification
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE__     =       logging.get_logger()






@dataclass
class     a_						:
  lowercase							=				42
  lowercase							=				field(default_factory=lowerCamelCase    )
  lowercase							=				field(default_factory=lowerCamelCase    )

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				int:



        """simple docstring"""






        UpperCamelCase      = len(list(m.modules()				)				) == 1 or isinstance(_SCREAMING_SNAKE_CASE      ,		nn.Convad				) or isinstance(_SCREAMING_SNAKE_CASE      ,		nn.BatchNormad				)
        if has_not_submodules:
              self.traced.append(_SCREAMING_SNAKE_CASE				)

  def __call__(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Any:



        """simple docstring"""






        for m in self.module.modules():
              self.handles.append(m.register_forward_hook(self._forward_hook				)				)
        self.module(_SCREAMING_SNAKE_CASE				)
        [x.remove() for x in self.handles]
        return self


  @property
  def   A__						(				self				)  ->				Tuple:



        """simple docstring"""






        return list(filter(lambda _SCREAMING_SNAKE_CASE				: len(list(x.state_dict().keys()				)				) > 0      ,		self.traced				)				)






@dataclass
class     a_						:
  lowercase							=				42
  lowercase							=				42
  lowercase							=				0
  lowercase							=				field(default_factory=lowerCamelCase    )
  lowercase							=				field(default_factory=lowerCamelCase    )

  def __call__(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				List[Any]:



        """simple docstring"""






        UpperCamelCase      = Tracker(self.dest				)(_SCREAMING_SNAKE_CASE				).parametrized
        UpperCamelCase      = Tracker(self.src				)(_SCREAMING_SNAKE_CASE				).parametrized

        UpperCamelCase      = list(filter(lambda _SCREAMING_SNAKE_CASE				: type(_SCREAMING_SNAKE_CASE				) not in self.src_skip      ,		_SCREAMING_SNAKE_CASE				)				)
        UpperCamelCase      = list(filter(lambda _SCREAMING_SNAKE_CASE				: type(_SCREAMING_SNAKE_CASE				) not in self.dest_skip      ,		_SCREAMING_SNAKE_CASE				)				)

        if len(_SCREAMING_SNAKE_CASE				) != len(_SCREAMING_SNAKE_CASE				):
              raise Exception(
                  F"Numbers of operations are different. Source module has {len(_SCREAMING_SNAKE_CASE				)} operations while"
                  F" destination module has {len(_SCREAMING_SNAKE_CASE				)}."				)

        for dest_m, src_m in zip(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				):
              dest_m.load_state_dict(src_m.state_dict()				)
              if self.verbose == 1:
                    print(F"Transfered from={src_m} to={dest_m}"				)







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase = True   )->			Optional[Any]:
      print(F"Converting {name}..."   )
      with torch.no_grad():
            UpperCamelCase      = timm.create_model(__UpperCamelCase		,						pretrained=__UpperCamelCase   ).eval()
            UpperCamelCase      = ResNetForImageClassification(__UpperCamelCase   ).eval()
            UpperCamelCase      = ModuleTransfer(src=__UpperCamelCase		,						dest=__UpperCamelCase   )
            UpperCamelCase      = torch.randn((1, 3, 224, 224)   )
            module_transfer(__UpperCamelCase   )

      assert torch.allclose(from_model(__UpperCamelCase   )		,						our_model(__UpperCamelCase   ).logits   ), "The model logits don't match the original one."

      UpperCamelCase      = F"resnet{'-'.join(name.split('resnet'   )   )}"
      print(__UpperCamelCase   )

      if push_to_hub:
            our_model.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name		,						commit_message="""Add model"""		,						use_temp_dir=__UpperCamelCase		,						)

            # we can use the convnext one
            UpperCamelCase      = AutoImageProcessor.from_pretrained("""facebook/convnext-base-224-22k-1k"""   )
            image_processor.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name		,						commit_message="""Add image processor"""		,						use_temp_dir=__UpperCamelCase		,						)

            print(F"Pushed {checkpoint_name}"   )







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase = None		,						__UpperCamelCase = True   )->			str:
      UpperCamelCase      = """imagenet-1k-id2label.json"""
      UpperCamelCase      = 1000
      UpperCamelCase      = (1, num_labels)

      UpperCamelCase      = """huggingface/label-files"""
      UpperCamelCase      = num_labels
      UpperCamelCase      = json.load(open(hf_hub_download(__UpperCamelCase		,						__UpperCamelCase		,						repo_type="""dataset"""   )		,						"""r"""   )   )
      UpperCamelCase      = {int(__UpperCamelCase   ): v for k, v in idalabel.items()}

      UpperCamelCase      = idalabel
      UpperCamelCase      = {v: k for k, v in idalabel.items()}

      UpperCamelCase      = partial(__UpperCamelCase		,						num_labels=__UpperCamelCase		,						idalabel=__UpperCamelCase		,						labelaid=__UpperCamelCase   )

      UpperCamelCase      = {
          """resnet18""": ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2]		,						hidden_sizes=[64, 128, 256, 512]		,						layer_type="""basic"""   ),
          """resnet26""": ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2]		,						hidden_sizes=[256, 512, 1024, 2048]		,						layer_type="""bottleneck"""   ),
          """resnet34""": ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3]		,						hidden_sizes=[64, 128, 256, 512]		,						layer_type="""basic"""   ),
          """resnet50""": ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3]		,						hidden_sizes=[256, 512, 1024, 2048]		,						layer_type="""bottleneck"""   ),
          """resnet101""": ImageNetPreTrainedConfig(
              depths=[3, 4, 23, 3]		,						hidden_sizes=[256, 512, 1024, 2048]		,						layer_type="""bottleneck"""   ),
          """resnet152""": ImageNetPreTrainedConfig(
              depths=[3, 8, 36, 3]		,						hidden_sizes=[256, 512, 1024, 2048]		,						layer_type="""bottleneck"""   ),
      }

      if model_name:
            convert_weight_and_push(__UpperCamelCase		,						names_to_config[model_name]		,						__UpperCamelCase		,						__UpperCamelCase   )
      else:
            for model_name, config in names_to_config.items():
                  convert_weight_and_push(__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   )
      return config, expected_shape


if __name__ == "__main__":
     SCREAMING_SNAKE_CASE__     =       argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default=None,
         type=str,
         help=(
             'The name of the model you wish to convert, it must be one of the supported resnet* architecture,'
             ' currently: resnet18,26,34,50,101,152. If `None`, all of them will the converted.'
         ),
     )
     parser.add_argument(
         '--pytorch_dump_folder_path',
         default=None,
         type=Path,
         required=True,
         help='Path to the output PyTorch model directory.',
     )
     parser.add_argument(
         '--push_to_hub',
         default=True,
         type=bool,
         required=False,
         help='If True, push model and image processor to the hub.',
     )

     SCREAMING_SNAKE_CASE__     =       parser.parse_args()
     SCREAMING_SNAKE_CASE__     =       args.pytorch_dump_folder_path
     pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
     convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)
