



import argparse
import gc
import json
import os

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.deepspeed import DummyOptim, DummyScheduler


__lowerCamelCase         =			16
__lowerCamelCase         =			32







def    _snake_case (       __snake_case    )     ->     str:






      '''simple docstring'''







      return int(x / 2**2_0    )


class   snake_case_ :




       """simple docstring"""






       def __enter__(     self):

             """simple docstring"""





             gc.collect()
             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
             UpperCAmelCase_  :	Optional[int]				      =		torch.cuda.memory_allocated()
             return self



       def __exit__(     self     ,*lowercase):

             """simple docstring"""





             gc.collect()
             torch.cuda.empty_cache()
             UpperCAmelCase_  :	Optional[Any]				      =		torch.cuda.memory_allocated()
             UpperCAmelCase_  :	Optional[Any]				      =		torch.cuda.max_memory_allocated()
             UpperCAmelCase_  :	int				      =		bamb(self.end - self.begin)
             UpperCAmelCase_  :	Tuple				      =		bamb(self.peak - self.begin)
             # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")







def    _snake_case (       __snake_case					,  __snake_case = 1_6					,  __snake_case = "bert-base-cased"					,  __snake_case = 3_2_0					,  __snake_case = 1_6_0					,  )     ->     Optional[int]:






      '''simple docstring'''







      UpperCAmelCase_  :	List[str]				      =		AutoTokenizer.from_pretrained(__snake_case    )
      UpperCAmelCase_  :	int				      =		load_dataset(
          "glue"					,  "mrpc"					,  split={"train": F"""train[:{n_train}]""", "validation": F"""validation[:{n_val}]"""}    )

      def tokenize_function(__snake_case    ):
            # max_length=None => use the model max length (it's actually the default)
            UpperCAmelCase_  :	int				      =		tokenizer(examples["sentence1"]					,  examples["sentence2"]					,  truncation=__snake_case					,  max_length=__snake_case    )
            return outputs

      # Apply the method we just defined to all the examples in all the splits of the dataset
      UpperCAmelCase_  :	Dict				      =		datasets.map(
          __snake_case					,  batched=__snake_case					,  remove_columns=["idx", "sentence1", "sentence2"]					,  load_from_cache_file=__snake_case    )

      # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
      # transformers library
      UpperCAmelCase_  :	Union[str, Any]				      =		tokenized_datasets.rename_column("label"					,  "labels"    )

      def collate_fn(__snake_case    ):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            if accelerator.distributed_type == DistributedType.TPU:
                  return tokenizer.pad(__snake_case					,  padding="max_length"					,  max_length=1_2_8					,  return_tensors="pt"    )
            return tokenizer.pad(__snake_case					,  padding="longest"					,  return_tensors="pt"    )

      # Instantiate dataloaders.
      UpperCAmelCase_  :	Optional[Any]				      =		DataLoader(
          tokenized_datasets["train"]					,  shuffle=__snake_case					,  collate_fn=__snake_case					,  batch_size=__snake_case    )
      UpperCAmelCase_  :	Union[str, Any]				      =		DataLoader(
          tokenized_datasets["validation"]					,  shuffle=__snake_case					,  collate_fn=__snake_case					,  batch_size=__snake_case    )

      return train_dataloader, eval_dataloader







def    _snake_case (       __snake_case					,  __snake_case    )     ->     Dict:






      '''simple docstring'''







      UpperCAmelCase_  :	str				      =		Accelerator()

      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      UpperCAmelCase_  :	Any				      =		config["lr"]
      UpperCAmelCase_  :	List[str]				      =		int(config["num_epochs"]    )
      UpperCAmelCase_  :	Optional[int]				      =		int(config["seed"]    )
      UpperCAmelCase_  :	Any				      =		int(config["batch_size"]    )
      UpperCAmelCase_  :	Union[str, Any]				      =		args.model_name_or_path

      set_seed(__snake_case    )
      UpperCAmelCase_      ,      UpperCAmelCase_  :	Optional[int]				      =		get_dataloaders(__snake_case					,  __snake_case					,  __snake_case					,  args.n_train					,  args.n_val    )

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      UpperCAmelCase_  :	Optional[int]				      =		AutoModelForSequenceClassification.from_pretrained(__snake_case					,  return_dict=__snake_case    )

      # Instantiate optimizer
      UpperCAmelCase_  :	Any				      =		(
          AdamW
          if accelerator.state.deepspeed_plugin is None
          or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
          else DummyOptim
      )
      UpperCAmelCase_  :	Any				      =		optimizer_cls(params=model.parameters()					,  lr=__snake_case    )

      if accelerator.state.deepspeed_plugin is not None:
            UpperCAmelCase_  :	Optional[int]				      =		accelerator.state.deepspeed_plugin.deepspeed_config[
                "gradient_accumulation_steps"
            ]
      else:
            UpperCAmelCase_  :	str				      =		1
      UpperCAmelCase_  :	Optional[int]				      =		(len(__snake_case    ) * num_epochs) // gradient_accumulation_steps

      # Instantiate scheduler
      if (
          accelerator.state.deepspeed_plugin is None
          or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
      ):
            UpperCAmelCase_  :	List[str]				      =		get_linear_schedule_with_warmup(
                optimizer=__snake_case					,  num_warmup_steps=0					,  num_training_steps=__snake_case					,  )
      else:
            UpperCAmelCase_  :	Union[str, Any]				      =		DummyScheduler(__snake_case					,  total_num_steps=__snake_case					,  warmup_num_steps=0    )

      # 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_  :	int				      =		accelerator.prepare(
          __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case    )

      # We need to keep track of how many total steps we have iterated over
      UpperCAmelCase_  :	Optional[int]				      =		0
      # We also need to keep track of the stating epoch so files are named properly
      UpperCAmelCase_  :	Optional[Any]				      =		0

      # Now we train the model
      UpperCAmelCase_  :	Optional[int]				      =		{}
      for epoch in range(__snake_case					,  __snake_case    ):
            with TorchTracemalloc() as tracemalloc:
                  model.train()
                  for step, batch in enumerate(__snake_case    ):
                        UpperCAmelCase_  :	Dict				      =		model(**__snake_case    )
                        UpperCAmelCase_  :	Optional[int]				      =		outputs.loss
                        UpperCAmelCase_  :	str				      =		loss / gradient_accumulation_steps
                        accelerator.backward(__snake_case    )
                        if step % gradient_accumulation_steps == 0:
                              optimizer.step()
                              lr_scheduler.step()
                              optimizer.zero_grad()

                        overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
            accelerator.print("Memory before entering the train : {}".format(bamb(tracemalloc.begin    )    )    )
            accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used    )    )
            accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked    )    )
            accelerator.print(
                "Total Peak Memory consumed during the train (max): {}".format(
                    tracemalloc.peaked + bamb(tracemalloc.begin    )    )    )
            UpperCAmelCase_  :	List[str]				      =		tracemalloc.peaked + bamb(tracemalloc.begin    )
            if args.peak_memory_upper_bound is not None:
                  assert (
                      train_total_peak_memory[F"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
                  ), "Peak memory usage exceeded the upper bound"

      accelerator.wait_for_everyone()
      if accelerator.is_main_process:
            with open(os.path.join(args.output_dir					,  "peak_memory_utilization.json"    )					,  "w"    ) as f:
                  json.dump(__snake_case					,  __snake_case    )







def    _snake_case (       )     ->     Union[str, Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	List[Any]				      =		argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."    )
      parser.add_argument(
          "--model_name_or_path"					,  type=__snake_case					,  default="bert-base-cased"					,  help="Path to pretrained model or model identifier from huggingface.co/models."					,  required=__snake_case					,  )
      parser.add_argument(
          "--output_dir"					,  type=__snake_case					,  default="."					,  help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."					,  )
      parser.add_argument(
          "--peak_memory_upper_bound"					,  type=__snake_case					,  default=__snake_case					,  help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value."					,  )
      parser.add_argument(
          "--n_train"					,  type=__snake_case					,  default=3_2_0					,  help="Number of training examples to use."					,  )
      parser.add_argument(
          "--n_val"					,  type=__snake_case					,  default=1_6_0					,  help="Number of validation examples to use."					,  )
      parser.add_argument(
          "--num_epochs"					,  type=__snake_case					,  default=1					,  help="Number of train epochs."					,  )
      UpperCAmelCase_  :	Optional[Any]				      =		parser.parse_args()
      UpperCAmelCase_  :	int				      =		{"lr": 2E-5, "num_epochs": args.num_epochs, "seed": 4_2, "batch_size": 1_6}
      training_function(__snake_case					,  __snake_case    )


if __name__ == "__main__":
   main()





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''microsoft/focalnet-tiny''': '''https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json''',
}


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""focalnet"""



       def __init__(     self     ,lowercase=224     ,lowercase=4     ,lowercase=3     ,lowercase=96     ,lowercase=False     ,lowercase=[192, 384, 768, 768]     ,lowercase=[2, 2, 6, 2]     ,lowercase=[2, 2, 2, 2]     ,lowercase=[3, 3, 3, 3]     ,lowercase="gelu"     ,lowercase=4.0     ,lowercase=0.0     ,lowercase=0.1     ,lowercase=False     ,lowercase=1E-4     ,lowercase=False     ,lowercase=False     ,lowercase=False     ,lowercase=0.02     ,lowercase=1E-5     ,lowercase=32     ,lowercase=None     ,lowercase=None     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(**lowercase)

             UpperCAmelCase_  :	int				      =		image_size
             UpperCAmelCase_  :	Optional[Any]				      =		patch_size
             UpperCAmelCase_  :	Dict				      =		num_channels
             UpperCAmelCase_  :	Dict				      =		embed_dim
             UpperCAmelCase_  :	Optional[int]				      =		use_conv_embed
             UpperCAmelCase_  :	int				      =		hidden_sizes
             UpperCAmelCase_  :	Optional[int]				      =		depths
             UpperCAmelCase_  :	Optional[Any]				      =		focal_levels
             UpperCAmelCase_  :	Tuple				      =		focal_windows
             UpperCAmelCase_  :	int				      =		hidden_act
             UpperCAmelCase_  :	Any				      =		mlp_ratio
             UpperCAmelCase_  :	List[Any]				      =		hidden_dropout_prob
             UpperCAmelCase_  :	Tuple				      =		drop_path_rate
             UpperCAmelCase_  :	List[str]				      =		use_layerscale
             UpperCAmelCase_  :	List[str]				      =		layerscale_value
             UpperCAmelCase_  :	List[str]				      =		use_post_layernorm
             UpperCAmelCase_  :	Dict				      =		use_post_layernorm_in_modulation
             UpperCAmelCase_  :	int				      =		normalize_modulator
             UpperCAmelCase_  :	Dict				      =		initializer_range
             UpperCAmelCase_  :	str				      =		layer_norm_eps
             UpperCAmelCase_  :	Union[str, Any]				      =		encoder_stride
             UpperCAmelCase_  :	str				      =		["stem"] + [F"""stage{idx}""" for idx in range(1     ,len(self.depths) + 1)]
             UpperCAmelCase_      ,      UpperCAmelCase_  :	Any				      =		get_aligned_output_features_output_indices(
                 out_features=lowercase     ,out_indices=lowercase     ,stage_names=self.stage_names)
