





"""simple docstring"""


import inspect
import unittest

from transformers import ResNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import ResNetBackbone, ResNetForImageClassification, ResNetModel
   from transformers.models.resnet.modeling_resnet import RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor



class   lowercase_					:

       def __init__(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE=3 ,       __SCREAMING_SNAKE_CASE=32 ,       __SCREAMING_SNAKE_CASE=3 ,       __SCREAMING_SNAKE_CASE=10 ,       __SCREAMING_SNAKE_CASE=[10, 20, 30, 40] ,       __SCREAMING_SNAKE_CASE=[1, 1, 2, 1] ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE=True ,       __SCREAMING_SNAKE_CASE="relu" ,       __SCREAMING_SNAKE_CASE=3 ,       __SCREAMING_SNAKE_CASE=None ,       __SCREAMING_SNAKE_CASE=["stage2", "stage3", "stage4"] ,       __SCREAMING_SNAKE_CASE=[2, 3, 4] ,       ):






        '''simple docstring'''

        snake_case__     				=					parent
        snake_case__     				=					batch_size
        snake_case__     				=					image_size
        snake_case__     				=					num_channels
        snake_case__     				=					embeddings_size
        snake_case__     				=					hidden_sizes
        snake_case__     				=					depths
        snake_case__     				=					is_training
        snake_case__     				=					use_labels
        snake_case__     				=					hidden_act
        snake_case__     				=					num_labels
        snake_case__     				=					scope
        snake_case__     				=					len(snake_case__  )
        snake_case__     				=					out_features
        snake_case__     				=					out_indices

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

        snake_case__     				=					None
        if self.use_labels:
         snake_case__     				=					ids_tensor([self.batch_size] ,       self.num_labels  )

        snake_case__     				=					self.get_config()

        return config, pixel_values, labels

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        return ResNetConfig(
            num_channels=self.num_channels ,       embeddings_size=self.embeddings_size ,       hidden_sizes=self.hidden_sizes ,       depths=self.depths ,       hidden_act=self.hidden_act ,       num_labels=self.num_labels ,       out_features=self.out_features ,       out_indices=self.out_indices ,       )

       def 		_SCREAMING_SNAKE_CASE(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  ):






        '''simple docstring'''

        snake_case__     				=					ResNetModel(config=snake_case__  )
        model.to(snake_case__  )
        model.eval()
        snake_case__     				=					model(snake_case__  )
        # expected last hidden states: B, C, H // 32, W // 32
        self.parent.assertEqual(
            result.last_hidden_state.shape ,       (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32) ,       )

       def 		_SCREAMING_SNAKE_CASE(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  ):






        '''simple docstring'''

        snake_case__     				=					self.num_labels
        snake_case__     				=					ResNetForImageClassification(snake_case__  )
        model.to(snake_case__  )
        model.eval()
        snake_case__     				=					model(snake_case__ ,       labels=snake_case__  )
        self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)  )

       def 		_SCREAMING_SNAKE_CASE(      self ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  ):






        '''simple docstring'''

        snake_case__     				=					ResNetBackbone(config=snake_case__  )
        model.to(snake_case__  )
        model.eval()
        snake_case__     				=					model(snake_case__  )

        # verify feature maps
        self.parent.assertEqual(len(result.feature_maps  ) ,       len(config.out_features  )  )
        self.parent.assertListEqual(list(result.feature_maps[0].shape  ) ,       [self.batch_size, self.hidden_sizes[1], 4, 4]  )

        # verify channels
        self.parent.assertEqual(len(model.channels  ) ,       len(config.out_features  )  )
        self.parent.assertListEqual(model.channels ,       config.hidden_sizes[1:]  )

        # verify backbone works with out_features=None
        snake_case__     				=					None
        snake_case__     				=					ResNetBackbone(config=snake_case__  )
        model.to(snake_case__  )
        model.eval()
        snake_case__     				=					model(snake_case__  )

        # verify feature maps
        self.parent.assertEqual(len(result.feature_maps  ) ,       1  )
        self.parent.assertListEqual(list(result.feature_maps[0].shape  ) ,       [self.batch_size, self.hidden_sizes[-1], 1, 1]  )

        # verify channels
        self.parent.assertEqual(len(model.channels  ) ,       1  )
        self.parent.assertListEqual(model.channels ,       [config.hidden_sizes[-1]]  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.prepare_config_and_inputs()
        snake_case__				, snake_case__				, snake_case__     				=					config_and_inputs
        snake_case__     				=					{"pixel_values": pixel_values}
        return config, inputs_dict



@require_torch
class   lowercase_					(  lowercase_					,lowercase_					,unittest.TestCase				):
       snake_case__					:				Dict									=  (
           (
               ResNetModel,
               ResNetForImageClassification,
               ResNetBackbone,
           )
           if is_torch_available()
           else ()
       )
       snake_case__					:				Optional[Any]									=  (
           {"""feature-extraction""": ResNetModel, """image-classification""": ResNetForImageClassification}
           if is_torch_available()
           else {}
       )

       snake_case__					:				Any									=  True
       snake_case__					:				Union[str, Any]									=  False
       snake_case__					:				Dict									=  False
       snake_case__					:				Optional[Any]									=  False
       snake_case__					:				int									=  False

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					ResNetModelTester(self  )
        snake_case__     				=					ConfigTester(self ,       config_class=snake_case__ ,       has_text_modality=snake_case__  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        return

       @unittest.skip(reason="ResNet does not use inputs_embeds"  )
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       @unittest.skip(reason="ResNet does not support input and output embeddings"  )
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__				, snake_case__     				=					self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         snake_case__     				=					model_class(snake_case__  )
         snake_case__     				=					inspect.signature(model.forward  )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         snake_case__     				=					[*signature.parameters.keys()]

         snake_case__     				=					["pixel_values"]
         self.assertListEqual(arg_names[:1] ,       snake_case__  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*snake_case__  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*snake_case__  )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__				, snake_case__     				=					self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         snake_case__     				=					model_class(config=snake_case__  )
         for name, module in model.named_modules():
          if isinstance(snake_case__ ,       (nn.BatchNormad, nn.GroupNorm)  ):
           self.assertTrue(
               torch.all(module.weight == 1  ) ,       msg=F"Parameter {name} of model {model_class} seems not properly initialized" ,       )
           self.assertTrue(
               torch.all(module.bias == 0  ) ,       msg=F"Parameter {name} of model {model_class} seems not properly initialized" ,       )

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        def check_hidden_states_output(__SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE ,       __SCREAMING_SNAKE_CASE  ):
         snake_case__     				=					model_class(snake_case__  )
         model.to(snake_case__  )
         model.eval()

         with torch.no_grad():
          snake_case__     				=					model(**self._prepare_for_class(snake_case__ ,       snake_case__  )  )

         snake_case__     				=					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

         snake_case__     				=					self.model_tester.num_stages
         self.assertEqual(len(snake_case__  ) ,       expected_num_stages + 1  )

         # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
         self.assertListEqual(
             list(hidden_states[0].shape[-2:]  ) ,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4] ,       )

        snake_case__				, snake_case__     				=					self.model_tester.prepare_config_and_inputs_for_common()
        snake_case__     				=					["basic", "bottleneck"]
        for model_class in self.all_model_classes:
         for layer_type in layers_type:
          snake_case__     				=					layer_type
          snake_case__     				=					True
          check_hidden_states_output(snake_case__ ,       snake_case__ ,       snake_case__  )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          snake_case__     				=					True

          check_hidden_states_output(snake_case__ ,       snake_case__ ,       snake_case__  )

       @unittest.skip(reason="ResNet does not use feedforward chunking"  )
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*snake_case__  )

       @slow
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        for model_name in RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case__     				=					ResNetModel.from_pretrained(snake_case__  )
         self.assertIsNotNone(snake_case__  )





def 				a__      (	)     ->	List[str]:

 '''simple docstring'''






 snake_case__     				=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
 return image



@require_torch
@require_vision
class   lowercase_					(  unittest.TestCase				):

       @cached_property
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        return (
            AutoImageProcessor.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]  )
            if is_vision_available()
            else None
        )

       @slow
       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					ResNetForImageClassification.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]  ).to(snake_case__  )

        snake_case__     				=					self.default_image_processor
        snake_case__     				=					prepare_img()
        snake_case__     				=					image_processor(images=snake_case__ ,       return_tensors="pt"  ).to(snake_case__  )

        # forward pass
        with torch.no_grad():
         snake_case__     				=					model(**snake_case__  )

        # verify the logits
        snake_case__     				=					torch.Size((1, 1000)  )
        self.assertEqual(outputs.logits.shape ,       snake_case__  )

        snake_case__     				=					torch.tensor([-11.1069, -9.7877, -8.3777]  ).to(snake_case__  )

        self.assertTrue(torch.allclose(outputs.logits[0, :3] ,       snake_case__ ,       atol=1e-4  )  )



@require_torch
class   lowercase_					(  lowercase_					,unittest.TestCase				):
       snake_case__					:				Tuple									=  (ResNetBackbone,) if is_torch_available() else ()
       snake_case__					:				List[str]									=  False
       snake_case__					:				Any									=  ResNetConfig

       def 		_SCREAMING_SNAKE_CASE(      self  ):






        '''simple docstring'''

        snake_case__     				=					ResNetModelTester(self  )









import torch

from accelerate import Accelerator, DistributedType

class        snake_case_				:







   """simple docstring"""



   def __enter__(     self:						int						)	->		Tuple:


       """simple docstring"""



       if self.enabled:
           _lowercase     			=						self.model.no_sync()
           self.model_sync_obj.__enter__()

       return self

   def __exit__(     self:						str ,							snake_case__:						Tuple ,							snake_case__:						int ,							snake_case__:						Tuple						)	->		Any:


       """simple docstring"""



       if self.enabled:
           # Average all models on exit
           self._sync_and_avg_model_params()
           self.model_sync_obj.__exit__(snake_case__ ,							snake_case__ ,							snake_case__						)

   def __init__(     self:						int ,							snake_case__:						Accelerator ,							snake_case__:						torch.nn.Module ,							snake_case__:						int ,							snake_case__:						bool = True						)	->		Optional[Any]:


       """simple docstring"""



       if accelerator.distributed_type not in [
           DistributedType.NO,
           DistributedType.MULTI_CPU,
           DistributedType.MULTI_GPU,
       ]:
           raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'						)
       _lowercase     			=						enabled and accelerator.distributed_type != DistributedType.NO
       _lowercase     			=						0
       if self.enabled:
           _lowercase     			=						accelerator
           _lowercase     			=						model
           _lowercase     			=						local_sgd_steps

   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		List[Any]:


       """simple docstring"""



       self.num_steps += 1
       if not self.enabled:
           return

       if self.num_steps % self.local_sgd_steps == 0:
           self._sync_and_avg_model_params()

   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[int]						)	->		Dict:


       """simple docstring"""



       self.accelerator.wait_for_everyone()
       with self.accelerator.autocast():
           for param in self.model.parameters():
               _lowercase     			=						self.accelerator.reduce(param.data ,							reduction='mean'						)
