

import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


lowerCAmelCase__       = logging.getLogger(__name__)
def      _lowerCAmelCase(	__A=2   ,       __A=3   ,       __A=16   ,       __A = 10   ,       __A = 2			):

  def get_dataset(__A			):
    UpperCAmelCase					      =  torch.randn(batch_size * n_batches   ,       1			)
    return TensorDataset(__A   ,       a * x + b + 0.1 * torch.randn(batch_size * n_batches   ,       1			)			)

  UpperCAmelCase					      =  get_dataset(__A			)
  UpperCAmelCase					      =  get_dataset(__A			)
  UpperCAmelCase					      =  DataLoader(__A   ,       shuffle=__A   ,       batch_size=__A   ,       num_workers=4			)
  UpperCAmelCase					      =  DataLoader(__A   ,       shuffle=__A   ,       batch_size=__A   ,       num_workers=4			)
  return (train_dataloader, valid_dataloader)
def      _lowerCAmelCase(	__A   ,       __A   ,       __A   ,       __A   ,       __A   ,       __A=None			):
  UpperCAmelCase					      =  []
  for epoch in range(__A			):
    # Train quickly
    model.train()
    for batch in dataloader:
      UpperCAmelCase      ,	UpperCAmelCase					      =  batch
      UpperCAmelCase					      =  model(__A			)
      UpperCAmelCase					      =  torch.nn.functional.mse_loss(__A   ,       __A			)
      accelerator.backward(__A			)
      optimizer.step()
      optimizer.zero_grad()
    rands.append(random.random()			)  # Introduce some randomness
    if scheduler is not None:
      scheduler.step()
  return rands






class 						__magic_name__    (       nn.Module    ):


 def __init__(						self		: Union[str, Any]    )    ->      List[Any]:
   super().__init__()
   UpperCAmelCase					      =  nn.Parameter(torch.randn(1    )    )
   UpperCAmelCase					      =  nn.Parameter(torch.randn(1    )    )


 def _UpperCamelCase	(						self		: int			,						lowerCAmelCase__		: Optional[int]    )    ->      Union[str, Any]:
   return x * self.a + self.b






class 						__magic_name__    (       unittest.TestCase    ):


 def _UpperCamelCase	(						self		: Any    )    ->      Optional[int]:
   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     UpperCAmelCase					      =  ProjectConfiguration(total_limit=1			,						project_dir=lowerCAmelCase__			,						automatic_checkpoint_naming=lowerCAmelCase__    )
     # Train baseline
     UpperCAmelCase					      =  Accelerator(project_config=lowerCAmelCase__    )
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save initial
     accelerator.save_state()

     # Save second state
     accelerator.save_state()
     self.assertEqual(len(os.listdir(accelerator.project_dir    )    )			,						1    )


 def _UpperCamelCase	(						self		: Optional[Any]    )    ->      Tuple:
   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     # Train baseline
     UpperCAmelCase					      =  Accelerator()
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save initial
     UpperCAmelCase					      =  os.path.join(lowerCAmelCase__			,						"initial"    )
     accelerator.save_state(lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     UpperCAmelCase					      =  train(3			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()

     # Train partially
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     UpperCAmelCase					      =  Accelerator()
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     accelerator.load_state(lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )

     UpperCAmelCase					      =  train(2			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save everything
     UpperCAmelCase					      =  os.path.join(lowerCAmelCase__			,						"checkpoint"    )
     accelerator.save_state(lowerCAmelCase__    )

     # Load everything back in and make sure all states work
     accelerator.load_state(lowerCAmelCase__    )
     test_rands += train(1			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      List[Any]:
   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     UpperCAmelCase					      =  ProjectConfiguration(automatic_checkpoint_naming=lowerCAmelCase__    )

     # Train baseline
     UpperCAmelCase					      =  Accelerator(project_dir=lowerCAmelCase__			,						project_config=lowerCAmelCase__    )
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save initial
     accelerator.save_state()
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     UpperCAmelCase					      =  train(3			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()

     # Train partially
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     UpperCAmelCase					      =  ProjectConfiguration(iteration=1			,						automatic_checkpoint_naming=lowerCAmelCase__    )
     UpperCAmelCase					      =  Accelerator(project_dir=lowerCAmelCase__			,						project_config=lowerCAmelCase__    )
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     accelerator.load_state(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_0"    )    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )

     UpperCAmelCase					      =  train(2			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save everything
     accelerator.save_state()

     # Load everything back in and make sure all states work
     accelerator.load_state(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_1"    )    )
     test_rands += train(1			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     ((UpperCAmelCase)      ,	(UpperCAmelCase))					      =  model.a.item(), model.b.item()
     UpperCAmelCase					      =  optimizer.state_dict()
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertEqual(lowerCAmelCase__			,						lowerCAmelCase__    )


 def _UpperCamelCase	(						self		: Tuple    )    ->      List[Any]:
   UpperCAmelCase					      =  torch.tensor([1, 2, 3]    )
   UpperCAmelCase					      =  torch.tensor([2, 3, 4]    )
   UpperCAmelCase					      =  DummyModel()
   UpperCAmelCase					      =  torch.optim.Adam(net.parameters()    )
   UpperCAmelCase					      =  Accelerator()
   with self.assertRaises(lowerCAmelCase__    ) as ve:
     accelerator.register_for_checkpointing(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
   UpperCAmelCase					      =  str(ve.exception    )
   self.assertTrue("Item at index 0" in message    )
   self.assertTrue("Item at index 1" in message    )
   self.assertFalse("Item at index 2" in message    )
   self.assertFalse("Item at index 3" in message    )


 def _UpperCamelCase	(						self		: Tuple    )    ->      int:
   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  torch.optim.Adam(params=model.parameters()			,						lr=1e-3    )
     UpperCAmelCase					      =  torch.optim.lr_scheduler.StepLR(lowerCAmelCase__			,						step_size=1			,						gamma=0.99    )
     UpperCAmelCase      ,	UpperCAmelCase					      =  dummy_dataloaders()
     UpperCAmelCase					      =  ProjectConfiguration(automatic_checkpoint_naming=lowerCAmelCase__    )
     # Train baseline
     UpperCAmelCase					      =  Accelerator(project_dir=lowerCAmelCase__			,						project_config=lowerCAmelCase__    )
     UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase					      =  accelerator.prepare(
         lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     # Save initial
     accelerator.save_state()
     UpperCAmelCase					      =  scheduler.state_dict()
     train(3			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__    )
     self.assertNotEqual(lowerCAmelCase__			,						scheduler.state_dict()    )

     # Load everything back in and make sure all states work
     accelerator.load_state(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_0"    )    )
     self.assertEqual(lowerCAmelCase__			,						scheduler.state_dict()    )


 def _UpperCamelCase	(						self		: List[str]    )    ->      List[Any]:
   with tempfile.TemporaryDirectory() as tmpdir:
     set_seed(4_2    )
     UpperCAmelCase					      =  DummyModel()
     UpperCAmelCase					      =  ProjectConfiguration(automatic_checkpoint_naming=lowerCAmelCase__			,						total_limit=2    )
     # Train baseline
     UpperCAmelCase					      =  Accelerator(project_dir=lowerCAmelCase__			,						project_config=lowerCAmelCase__    )
     UpperCAmelCase					      =  accelerator.prepare(lowerCAmelCase__    )
     # Save 3 states:
     for _ in range(1_1    ):
       accelerator.save_state()
     self.assertTrue(not os.path.exists(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_0"    )    )    )
     self.assertTrue(os.path.exists(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_9"    )    )    )
     self.assertTrue(os.path.exists(os.path.join(lowerCAmelCase__			,						"checkpoints"			,						"checkpoint_10"    )    )    )


 @require_cuda
 def _UpperCamelCase	(						self		: List[Any]    )    ->      Optional[int]:
   UpperCAmelCase					      =  ["torchrun", f"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__    )]
   execute_subprocess_async(lowerCAmelCase__			,						env=os.environ.copy()    )


if __name__ == "__main__":
      lowerCAmelCase__       = "/tmp/accelerate/state_checkpointing"
      lowerCAmelCase__       = DummyModel()
      lowerCAmelCase__       = torch.optim.Adam(params=model.parameters(), lr=1e-3)
      lowerCAmelCase__       = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9_9)
      lowerCAmelCase__, lowerCAmelCase__       = dummy_dataloaders()
      lowerCAmelCase__       = ProjectConfiguration(automatic_checkpoint_naming=True)
      # Train baseline
      lowerCAmelCase__       = Accelerator(project_dir=savedir, project_config=project_config, mixed_precision="no")
      if accelerator.process_index == 0:
            if os.path.exists(savedir):
                  shutil.rmtree(savedir)
            os.makedirs(savedir)
      lowerCAmelCase__, lowerCAmelCase__, lowerCAmelCase__, lowerCAmelCase__, lowerCAmelCase__       = accelerator.prepare(
          model, optimizer, train_dataloader, valid_dataloader, scheduler
      )
      lowerCAmelCase__, lowerCAmelCase__       = accelerator.prepare(model, optimizer)
      train(3, model, train_dataloader, optimizer, accelerator, scheduler)
      # Check that the intial optimizer is loaded on the GPU
      for group in optimizer.param_groups:
            lowerCAmelCase__       = group["params"][0].device
            break
      assert param_device.type == accelerator.device.type
      lowerCAmelCase__       = model.cpu()
      accelerator.wait_for_everyone()
      accelerator.save_state()
      accelerator.wait_for_everyone()

      # Check CPU state
      accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="cpu")
      for group in optimizer.param_groups:
            lowerCAmelCase__       = group["params"][0].device
            break
      assert (
          param_device.type == torch.device("cpu").type
      ), f"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

      # Check device state
      model.to(accelerator.device)
      accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="on_device")
      for group in optimizer.param_groups:
            lowerCAmelCase__       = group["params"][0].device
            break
      assert (
          param_device.type == accelerator.device.type
      ), f"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

      # Check error
      with pytest.raises(TypeError, match="Unsupported optimizer map location passed"):
            accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="invalid")
      accelerator.wait_for_everyone()
      if accelerator.process_index == 0:
            shutil.rmtree(savedir)
      accelerator.wait_for_everyone()






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__       = {
    "facebook/convnextv2-tiny-1k-224": "https://huggingface.co/facebook/convnextv2-tiny-1k-224/resolve/main/config.json",
}






class 						__magic_name__    (       _snake_case	,						_snake_case    ):
 UpperCAmelCase							    =						"""convnextv2"""


 def __init__(						self		: Optional[Any]			,						lowerCAmelCase__		: List[Any]=3			,						lowerCAmelCase__		: str=4			,						lowerCAmelCase__		: Dict=4			,						lowerCAmelCase__		: Optional[Any]=None			,						lowerCAmelCase__		: Optional[Any]=None			,						lowerCAmelCase__		: str="gelu"			,						lowerCAmelCase__		: Optional[int]=0.02			,						lowerCAmelCase__		: Dict=1e-1_2			,						lowerCAmelCase__		: str=0.0			,						lowerCAmelCase__		: str=2_2_4			,						lowerCAmelCase__		: int=None			,						lowerCAmelCase__		: List[Any]=None			,						**lowerCAmelCase__		: List[Any]			,						)    ->      List[Any]:
   super().__init__(**lowerCAmelCase__    )

   UpperCAmelCase					      =  num_channels
   UpperCAmelCase					      =  patch_size
   UpperCAmelCase					      =  num_stages
   UpperCAmelCase					      =  [9_6, 1_9_2, 3_8_4, 7_6_8] if hidden_sizes is None else hidden_sizes
   UpperCAmelCase					      =  [3, 3, 9, 3] if depths is None else depths
   UpperCAmelCase					      =  hidden_act
   UpperCAmelCase					      =  initializer_range
   UpperCAmelCase					      =  layer_norm_eps
   UpperCAmelCase					      =  drop_path_rate
   UpperCAmelCase					      =  image_size
   UpperCAmelCase					      =  ["stem"] + [f"stage{idx}" for idx in range(1			,						len(self.depths    ) + 1    )]
   UpperCAmelCase      ,	UpperCAmelCase					      =  get_aligned_output_features_output_indices(
       out_features=lowerCAmelCase__			,						out_indices=lowerCAmelCase__			,						stage_names=self.stage_names    )



