





'''simple docstring'''

import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation







def  a__    (    lowercase :					List[str]				)		->		Any:





     """simple docstring"""

     _UpperCamelCase       =	384
     _UpperCamelCase       =	7
     if "tiny" in model_name:
          _UpperCamelCase       =	96
          _UpperCamelCase       =	(2, 2, 6, 2)
          _UpperCamelCase       =	(3, 6, 12, 24)
     elif "small" in model_name:
          _UpperCamelCase       =	96
          _UpperCamelCase       =	(2, 2, 18, 2)
          _UpperCamelCase       =	(3, 6, 12, 24)
     elif "base" in model_name:
          _UpperCamelCase       =	128
          _UpperCamelCase       =	(2, 2, 18, 2)
          _UpperCamelCase       =	(4, 8, 16, 32)
          _UpperCamelCase       =	12
          _UpperCamelCase       =	512
     elif "large" in model_name:
          _UpperCamelCase       =	192
          _UpperCamelCase       =	(2, 2, 18, 2)
          _UpperCamelCase       =	(6, 12, 24, 48)
          _UpperCamelCase       =	12
          _UpperCamelCase       =	768

     # set label information
     _UpperCamelCase       =	150
     _UpperCamelCase       =	'''huggingface/label-files'''
     _UpperCamelCase       =	'''ade20k-id2label.json'''
     _UpperCamelCase       =	json.load(open(hf_hub_download(__A, __A, repo_type='''dataset'''				), '''r'''				)				)
     _UpperCamelCase       =	{int(__A				): v for k, v in idalabel.items()}
     _UpperCamelCase       =	{v: k for k, v in idalabel.items()}

     _UpperCamelCase       =	SwinConfig(
         embed_dim=__A, depths=__A, num_heads=__A, window_size=__A, out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4'''], )
     _UpperCamelCase       =	UperNetConfig(
         backbone_config=__A, auxiliary_in_channels=__A, num_labels=__A, idalabel=__A, labelaid=__A, )

     return config







def  a__    (    lowercase :					List[str]				)		->		Optional[Any]:





     """simple docstring"""

     _UpperCamelCase       =	[]

     # fmt: off
     # stem
     rename_keys.append(('''backbone.patch_embed.projection.weight''', '''backbone.embeddings.patch_embeddings.projection.weight''')				)
     rename_keys.append(('''backbone.patch_embed.projection.bias''', '''backbone.embeddings.patch_embeddings.projection.bias''')				)
     rename_keys.append(('''backbone.patch_embed.norm.weight''', '''backbone.embeddings.norm.weight''')				)
     rename_keys.append(('''backbone.patch_embed.norm.bias''', '''backbone.embeddings.norm.bias''')				)
     # stages
     for i in range(len(config.backbone_config.depths				)				):
          for j in range(config.backbone_config.depths[i]				):
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias""")				)

          if i < 3:
               rename_keys.append((F"""backbone.stages.{i}.downsample.reduction.weight""", F"""backbone.encoder.layers.{i}.downsample.reduction.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.downsample.norm.weight""", F"""backbone.encoder.layers.{i}.downsample.norm.weight""")				)
               rename_keys.append((F"""backbone.stages.{i}.downsample.norm.bias""", F"""backbone.encoder.layers.{i}.downsample.norm.bias""")				)
          rename_keys.append((F"""backbone.norm{i}.weight""", F"""backbone.hidden_states_norms.stage{i+1}.weight""")				)
          rename_keys.append((F"""backbone.norm{i}.bias""", F"""backbone.hidden_states_norms.stage{i+1}.bias""")				)

     # decode head
     rename_keys.extend(
         [
             ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
             ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
             ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
             ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
         ]				)
     # fmt: on

     return rename_keys







def  a__    (    lowercase :					Optional[int], lowercase :					Any, lowercase :					Tuple				)		->		Optional[Any]:





     """simple docstring"""

     _UpperCamelCase       =	dct.pop(__A				)
     _UpperCamelCase       =	val







def  a__    (    lowercase :					str, lowercase :					Tuple				)		->		List[str]:





     """simple docstring"""

     _UpperCamelCase       =	[int(backbone_config.embed_dim * 2**i				) for i in range(len(backbone_config.depths				)				)]
     for i in range(len(backbone_config.depths				)				):
          _UpperCamelCase       =	num_features[i]
          for j in range(backbone_config.depths[i]				):
               # fmt: off
               # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
               _UpperCamelCase       =	state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""				)
               _UpperCamelCase       =	state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias"""				)
               # next, add query, keys and values (in that order) to the state dict
               _UpperCamelCase       =	in_proj_weight[:dim, :]
               _UpperCamelCase       =	in_proj_bias[: dim]
               _UpperCamelCase       =	in_proj_weight[
                   dim : dim * 2, :
               ]
               _UpperCamelCase       =	in_proj_bias[
                   dim : dim * 2
               ]
               _UpperCamelCase       =	in_proj_weight[
                   -dim :, :
               ]
               _UpperCamelCase       =	in_proj_bias[-dim :]
               # fmt: on







def  a__    (    lowercase :					int				)		->		Optional[int]:





     """simple docstring"""

     _UpperCamelCase   ,    _UpperCamelCase       =	x.shape
     _UpperCamelCase       =	x.reshape(__A, 4, in_channel // 4				)
     _UpperCamelCase       =	x[:, [0, 2, 1, 3], :].transpose(1, 2				).reshape(__A, __A				)
     return x







def  a__    (    lowercase :					str				)		->		Tuple:





     """simple docstring"""

     _UpperCamelCase   ,    _UpperCamelCase       =	x.shape
     _UpperCamelCase       =	x.reshape(__A, in_channel // 4, 4				)
     _UpperCamelCase       =	x[:, :, [0, 2, 1, 3]].transpose(1, 2				).reshape(__A, __A				)

     return x







def  a__    (    lowercase :					Union[str, Any]				)		->		int:





     """simple docstring"""

     _UpperCamelCase       =	x.shape[0]
     _UpperCamelCase       =	x.reshape(4, in_channel // 4				)
     _UpperCamelCase       =	x[[0, 2, 1, 3], :].transpose(0, 1				).reshape(__A				)
     return x







def  a__    (    lowercase :					Optional[Any]				)		->		List[Any]:





     """simple docstring"""

     _UpperCamelCase       =	x.shape[0]
     _UpperCamelCase       =	x.reshape(in_channel // 4, 4				)
     _UpperCamelCase       =	x[:, [0, 2, 1, 3]].transpose(0, 1				).reshape(__A				)
     return x







def  a__    (    lowercase :					Any, lowercase :					Dict, lowercase :					Optional[int]				)		->		List[Any]:





     """simple docstring"""

     _UpperCamelCase       =	{
         '''upernet-swin-tiny''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth''',
         '''upernet-swin-small''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth''',
         '''upernet-swin-base''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth''',
         '''upernet-swin-large''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth''',
     }
     _UpperCamelCase       =	model_name_to_url[model_name]
     _UpperCamelCase       =	torch.hub.load_state_dict_from_url(__A, map_location='''cpu''', file_name=__A				)[
         '''state_dict'''
     ]

     for name, param in state_dict.items():
          print(__A, param.shape				)

     _UpperCamelCase       =	get_upernet_config(__A				)
     _UpperCamelCase       =	UperNetForSemanticSegmentation(__A				)
     model.eval()

     # replace "bn" => "batch_norm"
     for key in state_dict.copy().keys():
          _UpperCamelCase       =	state_dict.pop(__A				)
          if "bn" in key:
               _UpperCamelCase       =	key.replace('''bn''', '''batch_norm'''				)
          _UpperCamelCase       =	val

     # rename keys
     _UpperCamelCase       =	create_rename_keys(__A				)
     for src, dest in rename_keys:
          rename_key(__A, __A, __A				)
     read_in_q_k_v(__A, config.backbone_config				)

     # fix downsample parameters
     for key, value in state_dict.items():
          if "downsample" in key:
               if "reduction" in key:
                    _UpperCamelCase       =	reverse_correct_unfold_reduction_order(__A				)
               if "norm" in key:
                    _UpperCamelCase       =	reverse_correct_unfold_norm_order(__A				)

     model.load_state_dict(__A				)

     # verify on image
     _UpperCamelCase       =	'''https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'''
     _UpperCamelCase       =	Image.open(requests.get(__A, stream=__A				).raw				).convert('''RGB'''				)

     _UpperCamelCase       =	SegformerImageProcessor()
     _UpperCamelCase       =	processor(__A, return_tensors='''pt'''				).pixel_values

     with torch.no_grad():
          _UpperCamelCase       =	model(__A				)
          _UpperCamelCase       =	outputs.logits

     print(logits.shape				)
     print('''First values of logits:''', logits[0, 0, :3, :3]				)
     # assert values
     if model_name == "upernet-swin-tiny":
          _UpperCamelCase       =	torch.tensor(
              [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]				)
     elif model_name == "upernet-swin-small":
          _UpperCamelCase       =	torch.tensor(
              [[-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.0_9_0_8, -7.0_9_0_8, -6.8_5_3_4]]				)
     elif model_name == "upernet-swin-base":
          _UpperCamelCase       =	torch.tensor(
              [[-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.4_7_6_3, -6.4_7_6_3, -6.3_2_5_4]]				)
     elif model_name == "upernet-swin-large":
          _UpperCamelCase       =	torch.tensor(
              [[-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.4_0_4_4, -7.4_0_4_4, -7.2_5_8_6]]				)
     print('''Logits:''', outputs.logits[0, 0, :3, :3]				)
     assert torch.allclose(outputs.logits[0, 0, :3, :3], __A, atol=1e-4				)
     print('''Looks ok!'''				)

     if pytorch_dump_folder_path is not None:
          print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""				)
          model.save_pretrained(__A				)
          print(F"""Saving processor to {pytorch_dump_folder_path}"""				)
          processor.save_pretrained(__A				)

     if push_to_hub:
          print(F"""Pushing model and processor for {model_name} to hub"""				)
          model.push_to_hub(F"""openmmlab/{model_name}"""				)
          processor.push_to_hub(F"""openmmlab/{model_name}"""				)


if __name__ == "__main__":
  lowercase__						:      List[Any]	=    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--model_name',
      default='upernet-swin-tiny',
      type=str,
      choices=[F"""upernet-swin-{size}""" for size in ['tiny', 'small', 'base', 'large']],
      help='Name of the Swin + UperNet model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )
  parser.add_argument(
      '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
  )

  lowercase__						:      List[str]	=    parser.parse_args()
  convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''

import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowercase__						:      str	=    16
lowercase__						:      int	=    32







def  a__    (    lowercase :					Accelerator, lowercase :					int = 16				)		->		List[str]:





     """simple docstring"""

     _UpperCamelCase       =	AutoTokenizer.from_pretrained('''bert-base-cased'''				)
     _UpperCamelCase       =	load_dataset('''glue''', '''mrpc'''				)

     def tokenize_function(lowercase :					Union[str, Any]				):
          # max_length=None => use the model max length (it's actually the default)
          _UpperCamelCase       =	tokenizer(examples['''sentence1'''], examples['''sentence2'''], truncation=lowercase, max_length=lowercase				)
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          _UpperCamelCase       =	datasets.map(
              lowercase, batched=lowercase, remove_columns=['''idx''', '''sentence1''', '''sentence2'''], )

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     _UpperCamelCase       =	tokenized_datasets.rename_column('''label''', '''labels'''				)

     def collate_fn(lowercase :					Any				):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          _UpperCamelCase       =	128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               _UpperCamelCase       =	16
          elif accelerator.mixed_precision != "no":
               _UpperCamelCase       =	8
          else:
               _UpperCamelCase       =	None

          return tokenizer.pad(
              lowercase, padding='''longest''', max_length=lowercase, pad_to_multiple_of=lowercase, return_tensors='''pt''', )

     # Instantiate dataloaders.
     _UpperCamelCase       =	DataLoader(
         tokenized_datasets['''train'''], shuffle=lowercase, collate_fn=lowercase, batch_size=lowercase				)
     _UpperCamelCase       =	DataLoader(
         tokenized_datasets['''validation'''], shuffle=lowercase, collate_fn=lowercase, batch_size=lowercase				)

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   lowercase__						:      str	=    mocked_dataloaders  # noqa: F811







def  a__    (    lowercase :					List[Any], lowercase :					List[str]				)		->		Any:





     """simple docstring"""

     if os.environ.get('''TESTING_MOCKED_DATALOADERS''', lowercase				) == "1":
          _UpperCamelCase       =	2
     # Initialize accelerator
     _UpperCamelCase       =	Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision				)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     _UpperCamelCase       =	config['''lr''']
     _UpperCamelCase       =	int(config['''num_epochs''']				)
     _UpperCamelCase       =	int(config['''seed''']				)
     _UpperCamelCase       =	int(config['''batch_size''']				)

     _UpperCamelCase       =	evaluate.load('''glue''', '''mrpc'''				)

     # New Code #
     # We now can define an inner training loop function. It should take a batch size as the only parameter,
     # and build the dataloaders in there.
     # It also gets our decorator
     @find_executable_batch_size(starting_batch_size=lowercase				)
     def inner_training_loop(lowercase :					List[Any]				):
          # And now just move everything below under this function
          # We need to bring in the Accelerator object from earlier
          nonlocal accelerator
          # And reset all of its attributes that could hold onto any memory:
          accelerator.free_memory()

          # Then we can declare the model, optimizer, and everything else:
          set_seed(lowercase				)

          # Instantiate the model (we build the model here so that the seed also control new weights initialization)
          _UpperCamelCase       =	AutoModelForSequenceClassification.from_pretrained('''bert-base-cased''', return_dict=lowercase				)

          # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
          # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
          # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
          _UpperCamelCase       =	model.to(accelerator.device				)

          # Instantiate optimizer
          _UpperCamelCase       =	AdamW(params=model.parameters(), lr=lowercase				)
          _UpperCamelCase   ,    _UpperCamelCase       =	get_dataloaders(lowercase, lowercase				)

          # Instantiate scheduler
          _UpperCamelCase       =	get_linear_schedule_with_warmup(
              optimizer=lowercase, num_warmup_steps=100, num_training_steps=(len(lowercase				) * num_epochs), )

          # Prepare everything
          # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
          # prepare method.
          _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase       =	accelerator.prepare(
              lowercase, lowercase, lowercase, lowercase, lowercase				)

          # Now we train the model
          for epoch in range(lowercase				):
               model.train()
               for step, batch in enumerate(lowercase				):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device				)
                    _UpperCamelCase       =	model(**lowercase				)
                    _UpperCamelCase       =	outputs.loss
                    accelerator.backward(lowercase				)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               model.eval()
               for step, batch in enumerate(lowercase				):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device				)
                    with torch.no_grad():
                         _UpperCamelCase       =	model(**lowercase				)
                    _UpperCamelCase       =	outputs.logits.argmax(dim=-1				)
                    _UpperCamelCase   ,    _UpperCamelCase       =	accelerator.gather_for_metrics((predictions, batch['''labels'''])				)
                    metric.add_batch(
                        predictions=lowercase, references=lowercase, )

               _UpperCamelCase       =	metric.compute()
               # Use accelerator.print to print only on the main process.
               accelerator.print(F"""epoch {epoch}:""", lowercase				)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
     inner_training_loop()







def  a__    (    )		->		str:





     """simple docstring"""

     _UpperCamelCase       =	argparse.ArgumentParser(description='''Simple example of training script.'''				)
     parser.add_argument(
         '''--mixed_precision''', type=lowercase, default=lowercase, choices=['''no''', '''fp16''', '''bf16''', '''fp8'''], help='''Whether to use mixed precision. Choose'''
         '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
         '''and an Nvidia Ampere GPU.''', )
     parser.add_argument('''--cpu''', action='''store_true''', help='''If passed, will train on the CPU.'''				)
     _UpperCamelCase       =	parser.parse_args()
     _UpperCamelCase       =	{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
     training_function(lowercase, lowercase				)


if __name__ == "__main__":
   main()



