






import argparse
import os
import re

import numpy as np
import PIL
import torch
from timm import create_model
from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator


def        _a  (     UpperCamelCase_					:				int   ) ->   Optional[Any]:
						"""simple docstring"""







						lowerCAmelCase__       =						fname.split(os.path.sep   )[-1]
						return re.search(R"^(.*)_\d+\.jpg$"    ,   UpperCamelCase_   ).groups()[0]


class      lowercase__     (			_UpperCAmelCase	):

	def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						file_names
							lowerCAmelCase__       =						image_transform
							lowerCAmelCase__       =						label_to_id

	def __len__(							self       )->					Optional[Any]:







							'''simple docstring'''





							return len(self.file_names       )


	def __getitem__(							self		,    __UpperCAmelCase       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						self.file_names[idx]
							lowerCAmelCase__       =						PIL.Image.open(__UpperCAmelCase       )
							lowerCAmelCase__       =						raw_image.convert("RGB"       )
							if self.image_transform is not None:
													lowerCAmelCase__       =						self.image_transform(__UpperCAmelCase       )
							lowerCAmelCase__       =						extract_label(__UpperCAmelCase       )
							if self.label_to_id is not None:
													lowerCAmelCase__       =						self.label_to_id[label]
							return {"image": image, "label": label}





def        _a  (     UpperCamelCase_					:				List[Any]    ,   UpperCamelCase_					:				Union[str, Any]   ) ->   Optional[int]:
						"""simple docstring"""







						if args.with_tracking:
												lowerCAmelCase__       =						Accelerator(
												    cpu=args.cpu    ,   mixed_precision=args.mixed_precision    ,   log_with="all"    ,   project_dir=args.project_dir   )
						else:
												lowerCAmelCase__       =						Accelerator(cpu=args.cpu    ,   mixed_precision=args.mixed_precision   )

						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						lowerCAmelCase__       =						config["lr"]
						lowerCAmelCase__       =						int(config["num_epochs"]   )
						lowerCAmelCase__       =						int(config["seed"]   )
						lowerCAmelCase__       =						int(config["batch_size"]   )
						lowerCAmelCase__       =						config["image_size"]
						if not isinstance(UpperCamelCase_    ,   (list, tuple)   ):
												lowerCAmelCase__       =						(image_size, image_size)

						# Parse out whether we are saving every epoch or after a certain number of batches
						if hasattr(args.checkpointing_steps    ,   "isdigit"   ):
												if args.checkpointing_steps == "epoch":
																		lowerCAmelCase__       =						args.checkpointing_steps
												elif args.checkpointing_steps.isdigit():
																		lowerCAmelCase__       =						int(args.checkpointing_steps   )
												else:
																		raise ValueError(
																		    F"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."   )
						else:
												lowerCAmelCase__       =						None

						# We need to initialize the trackers we use, and also store our configuration
						if args.with_tracking:
												lowerCAmelCase__       =						os.path.split(UpperCamelCase_   )[-1].split("."   )[0]
												accelerator.init_trackers(UpperCamelCase_    ,   UpperCamelCase_   )

						# Grab all the image filenames
						lowerCAmelCase__       =						[os.path.join(args.data_dir    ,   UpperCamelCase_   ) for fname in os.listdir(args.data_dir   ) if fname.endswith(".jpg"   )]

						# Build the label correspondences
						lowerCAmelCase__       =						[extract_label(UpperCamelCase_   ) for fname in file_names]
						lowerCAmelCase__       =						list(set(UpperCamelCase_   )   )
						id_to_label.sort()
						lowerCAmelCase__       =						{lbl: i for i, lbl in enumerate(UpperCamelCase_   )}

						# Set the seed before splitting the data.
						np.random.seed(UpperCamelCase_   )
						torch.manual_seed(UpperCamelCase_   )
						torch.cuda.manual_seed_all(UpperCamelCase_   )

						# Split our filenames between train and validation
						lowerCAmelCase__       =						np.random.permutation(len(UpperCamelCase_   )   )
						lowerCAmelCase__       =						int(0.8 * len(UpperCamelCase_   )   )
						lowerCAmelCase__       =						random_perm[:cut]
						lowerCAmelCase__       =						random_perm[cut:]

						# For training we use a simple RandomResizedCrop
						lowerCAmelCase__       =						Compose([RandomResizedCrop(UpperCamelCase_    ,   scale=(0.5, 1.0)   ), ToTensor()]   )
						lowerCAmelCase__       =						PetsDataset(
						    [file_names[i] for i in train_split]    ,   image_transform=UpperCamelCase_    ,   label_to_id=UpperCamelCase_   )

						# For evaluation, we use a deterministic Resize
						lowerCAmelCase__       =						Compose([Resize(UpperCamelCase_   ), ToTensor()]   )
						lowerCAmelCase__       =						PetsDataset([file_names[i] for i in eval_split]    ,   image_transform=UpperCamelCase_    ,   label_to_id=UpperCamelCase_   )

						# Instantiate dataloaders.
						lowerCAmelCase__       =						DataLoader(UpperCamelCase_    ,   shuffle=UpperCamelCase_    ,   batch_size=UpperCamelCase_    ,   num_workers=4   )
						lowerCAmelCase__       =						DataLoader(UpperCamelCase_    ,   shuffle=UpperCamelCase_    ,   batch_size=UpperCamelCase_    ,   num_workers=4   )

						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						lowerCAmelCase__       =						create_model("resnet50d"    ,   pretrained=UpperCamelCase_    ,   num_classes=len(UpperCamelCase_   )   )

						# 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).
						lowerCAmelCase__       =						model.to(accelerator.device   )

						# Freezing the base model
						for param in model.parameters():
												lowerCAmelCase__       =						False
						for param in model.get_classifier().parameters():
												lowerCAmelCase__       =						True

						# We normalize the batches of images to be a bit faster.
						lowerCAmelCase__       =						torch.tensor(model.default_cfg["mean"]   )[None, :, None, None].to(accelerator.device   )
						lowerCAmelCase__       =						torch.tensor(model.default_cfg["std"]   )[None, :, None, None].to(accelerator.device   )

						# Instantiate optimizer
						lowerCAmelCase__       =						torch.optim.Adam(params=model.parameters()    ,   lr=lr / 25   )

						# Instantiate learning rate scheduler
						lowerCAmelCase__       =						OneCycleLR(optimizer=UpperCamelCase_    ,   max_lr=UpperCamelCase_    ,   epochs=UpperCamelCase_    ,   steps_per_epoch=len(UpperCamelCase_   )   )

						# 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.
						lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						accelerator.prepare(
						    UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_   )
						# We need to keep track of how many total steps we have iterated over
						lowerCAmelCase__       =						0
						# We also need to keep track of the starting epoch so files are named properly
						lowerCAmelCase__       =						0

						# Potentially load in the weights and states from a previous save
						if args.resume_from_checkpoint:
												if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
																		accelerator.print(F"Resumed from checkpoint: {args.resume_from_checkpoint}"   )
																		accelerator.load_state(args.resume_from_checkpoint   )
																		lowerCAmelCase__       =						os.path.basename(args.resume_from_checkpoint   )
												else:
																		# Get the most recent checkpoint
																		lowerCAmelCase__       =						[f.name for f in os.scandir(os.getcwd()   ) if f.is_dir()]
																		dirs.sort(key=os.path.getctime   )
																		lowerCAmelCase__       =						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
												# Extract `epoch_{i}` or `step_{i}`
												lowerCAmelCase__       =						os.path.splitext(UpperCamelCase_   )[0]

												if "epoch" in training_difference:
																		lowerCAmelCase__       =						int(training_difference.replace("epoch_"    ,   ""   )   ) + 1
																		lowerCAmelCase__       =						None
												else:
																		lowerCAmelCase__       =						int(training_difference.replace("step_"    ,   ""   )   )
																		lowerCAmelCase__       =						resume_step // len(UpperCamelCase_   )
																		resume_step -= starting_epoch * len(UpperCamelCase_   )

    # Now we train the model
						for epoch in range(UpperCamelCase_    ,   UpperCamelCase_   ):
												model.train()
												if args.with_tracking:
																		lowerCAmelCase__       =						0
												if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
																		# We need to skip steps until we reach the resumed step
																		lowerCAmelCase__       =						accelerator.skip_first_batches(UpperCamelCase_    ,   UpperCamelCase_   )
																		overall_step += resume_step
												else:
																		# After the first iteration though, we need to go back to the original dataloader
																		lowerCAmelCase__       =						train_dataloader
												for batch in active_dataloader:
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		lowerCAmelCase__       =						{k: v.to(accelerator.device   ) for k, v in batch.items()}
																		lowerCAmelCase__       =						(batch["image"] - mean) / std
																		lowerCAmelCase__       =						model(UpperCamelCase_   )
																		lowerCAmelCase__       =						torch.nn.functional.cross_entropy(UpperCamelCase_    ,   batch["label"]   )
																		# We keep track of the loss at each epoch
																		if args.with_tracking:
																								total_loss += loss.detach().float()
																		accelerator.backward(UpperCamelCase_   )
																		optimizer.step()
																		lr_scheduler.step()
																		optimizer.zero_grad()
																		overall_step += 1
																		if isinstance(UpperCamelCase_    ,   UpperCamelCase_   ):
																								lowerCAmelCase__       =						F"step_{overall_step}"
																								if overall_step % checkpointing_steps == 0:
																														if args.output_dir is not None:
																																				lowerCAmelCase__       =						os.path.join(args.output_dir    ,   UpperCamelCase_   )
																														accelerator.save_state(UpperCamelCase_   )
												model.eval()
												lowerCAmelCase__       =						0
												lowerCAmelCase__       =						0
												for step, batch in enumerate(UpperCamelCase_   ):
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		lowerCAmelCase__       =						{k: v.to(accelerator.device   ) for k, v in batch.items()}
																		lowerCAmelCase__       =						(batch["image"] - mean) / std
																		with torch.no_grad():
																								lowerCAmelCase__       =						model(UpperCamelCase_   )
																		lowerCAmelCase__       =						outputs.argmax(dim=-1   )
																		lowerCAmelCase__      ,							lowerCAmelCase__       =						accelerator.gather_for_metrics((predictions, batch["label"])   )
																		lowerCAmelCase__       =						predictions == references
																		num_elems += accurate_preds.shape[0]
																		accurate += accurate_preds.long().sum()

												lowerCAmelCase__       =						accurate.item() / num_elems
												# Use accelerator.print to print only on the main process.
												accelerator.print(F"epoch {epoch}: {100 * eval_metric:.2f}"   )
												if args.with_tracking:
																		accelerator.log(
																		    {
																		        "accuracy": 100 * eval_metric,
																		        "train_loss": total_loss.item() / len(UpperCamelCase_   ),
																		        "epoch": epoch,
																		    }    ,   step=UpperCamelCase_    ,   )
												if checkpointing_steps == "epoch":
																		lowerCAmelCase__       =						F"epoch_{epoch}"
																		if args.output_dir is not None:
																								lowerCAmelCase__       =						os.path.join(args.output_dir    ,   UpperCamelCase_   )
																		accelerator.save_state(UpperCamelCase_   )

						if args.with_tracking:
												accelerator.end_training()


def        _a  (     ) ->   Tuple:
						"""simple docstring"""







						lowerCAmelCase__       =						argparse.ArgumentParser(description="Simple example of training script."   )
						parser.add_argument("--data_dir"    ,   required=UpperCamelCase_    ,   help="The data folder on disk."   )
						parser.add_argument("--fp16"    ,   action="store_true"    ,   help="If passed, will use FP16 training."   )
						parser.add_argument(
						    "--mixed_precision"    ,   type=UpperCamelCase_    ,   default=UpperCamelCase_    ,   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."   )
						parser.add_argument(
						    "--checkpointing_steps"    ,   type=UpperCamelCase_    ,   default=UpperCamelCase_    ,   help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."    ,   )
						parser.add_argument(
						    "--output_dir"    ,   type=UpperCamelCase_    ,   default="."    ,   help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."    ,   )
						parser.add_argument(
						    "--resume_from_checkpoint"    ,   type=UpperCamelCase_    ,   default=UpperCamelCase_    ,   help="If the training should continue from a checkpoint folder."    ,   )
						parser.add_argument(
						    "--with_tracking"    ,   action="store_true"    ,   help="Whether to load in all available experiment trackers from the environment and use them for logging."    ,   )
						parser.add_argument(
						    "--project_dir"    ,   type=UpperCamelCase_    ,   default="logs"    ,   help="Location on where to store experiment tracking logs` and relevent project information"    ,   )
						lowerCAmelCase__       =						parser.parse_args()
						lowerCAmelCase__       =						{"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
						training_function(UpperCamelCase_    ,   UpperCamelCase_   )


if __name__ == "__main__":
	main()












import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 transformers import (
	    MODEL_MAPPING,
	    SegformerForImageClassification,
	    SegformerForSemanticSegmentation,
	    SegformerModel,
	)
	from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import SegformerImageProcessor


class      lowercase__     (			_UpperCAmelCase	):

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.config_class(**self.inputs_dict       )
							self.parent.assertTrue(hasattr(__UpperCAmelCase		,    "hidden_sizes"       )       )
							self.parent.assertTrue(hasattr(__UpperCAmelCase		,    "num_attention_heads"       )       )
							self.parent.assertTrue(hasattr(__UpperCAmelCase		,    "num_encoder_blocks"       )       )





class      lowercase__     :

	def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=13		,    __UpperCAmelCase=64		,    __UpperCAmelCase=3		,    __UpperCAmelCase=4		,    __UpperCAmelCase=[2, 2, 2, 2]		,    __UpperCAmelCase=[8, 4, 2, 1]		,    __UpperCAmelCase=[16, 32, 64, 128]		,    __UpperCAmelCase=[1, 4, 8, 16]		,    __UpperCAmelCase=[1, 2, 4, 8]		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase="gelu"		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.02		,    __UpperCAmelCase=3		,    __UpperCAmelCase=None		,    )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						parent
							lowerCAmelCase__       =						batch_size
							lowerCAmelCase__       =						image_size
							lowerCAmelCase__       =						num_channels
							lowerCAmelCase__       =						num_encoder_blocks
							lowerCAmelCase__       =						sr_ratios
							lowerCAmelCase__       =						depths
							lowerCAmelCase__       =						hidden_sizes
							lowerCAmelCase__       =						downsampling_rates
							lowerCAmelCase__       =						num_attention_heads
							lowerCAmelCase__       =						is_training
							lowerCAmelCase__       =						use_labels
							lowerCAmelCase__       =						hidden_act
							lowerCAmelCase__       =						hidden_dropout_prob
							lowerCAmelCase__       =						attention_probs_dropout_prob
							lowerCAmelCase__       =						initializer_range
							lowerCAmelCase__       =						num_labels
							lowerCAmelCase__       =						scope

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





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

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

							lowerCAmelCase__       =						self.get_config()
							return config, pixel_values, labels

	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							return SegformerConfig(
							    image_size=self.image_size		,    num_channels=self.num_channels		,    num_encoder_blocks=self.num_encoder_blocks		,    depths=self.depths		,    hidden_sizes=self.hidden_sizes		,    num_attention_heads=self.num_attention_heads		,    hidden_act=self.hidden_act		,    hidden_dropout_prob=self.hidden_dropout_prob		,    attention_probs_dropout_prob=self.attention_probs_dropout_prob		,    initializer_range=self.initializer_range		,    )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__       =						SegformerModel(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase       )
							lowerCAmelCase__       =						lowerCAmelCase__       =						self.image_size // (self.downsampling_rates[-1] * 2)
							self.parent.assertEqual(
							    result.last_hidden_state.shape		,    (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						self.num_labels
							lowerCAmelCase__       =						SegformerForSemanticSegmentation(__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase       )
							self.parent.assertEqual(
							    result.logits.shape		,    (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    labels=__UpperCAmelCase       )
							self.parent.assertEqual(
							    result.logits.shape		,    (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
							self.parent.assertGreater(result.loss		,    0.0       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						1
							lowerCAmelCase__       =						SegformerForSemanticSegmentation(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						torch.randint(0		,    1		,    (self.batch_size, self.image_size, self.image_size)       ).to(__UpperCAmelCase       )
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    labels=__UpperCAmelCase       )
							self.parent.assertGreater(result.loss		,    0.0       )


	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





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





@require_torch
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase	):
	a_         =(
	    (
	        SegformerModel,
	        SegformerForSemanticSegmentation,
	        SegformerForImageClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	a_         =(
	    {
	        """feature-extraction""": SegformerModel,
	        """image-classification""": SegformerForImageClassification,
	        """image-segmentation""": SegformerForSemanticSegmentation,
	    }
	    if is_torch_available()
	    else {}
	)

	a_         =True
	a_         =False
	a_         =False
	a_         =False

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						SegformerModelTester(self       )
							lowerCAmelCase__       =						SegformerConfigTester(self		,    config_class=__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							self.config_tester.run_common_tests()

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_binary_image_segmentation(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_segmentation(*__UpperCAmelCase       )

	@unittest.skip("SegFormer does not use inputs_embeds"       )
	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							pass

	@unittest.skip("SegFormer does not have get_input_embeddings method and get_output_embeddings methods"       )
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							pass

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()

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

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

	def 		UpperCAmelCase  (							self       )->					str:







							'''simple docstring'''





							lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()
							lowerCAmelCase__       =						True

							for model_class in self.all_model_classes:
													lowerCAmelCase__       =						True
													lowerCAmelCase__       =						False
													lowerCAmelCase__       =						True
													lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
													model.to(__UpperCAmelCase       )
													model.eval()
													with torch.no_grad():
																			lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )
													lowerCAmelCase__       =						outputs.attentions

													lowerCAmelCase__       =						sum(self.model_tester.depths       )
													self.assertEqual(len(__UpperCAmelCase       )		,    __UpperCAmelCase       )

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													lowerCAmelCase__       =						True
													lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
													model.to(__UpperCAmelCase       )
													model.eval()
													with torch.no_grad():
																			lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )
													lowerCAmelCase__       =						outputs.attentions

													self.assertEqual(len(__UpperCAmelCase       )		,    __UpperCAmelCase       )

													# verify the first attentions (first block, first layer)
													lowerCAmelCase__       =						(self.model_tester.image_size // 4) ** 2
													lowerCAmelCase__       =						(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
													self.assertListEqual(
													    list(attentions[0].shape[-3:]       )		,    [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]		,    )

													# verify the last attentions (last block, last layer)
													lowerCAmelCase__       =						(self.model_tester.image_size // 32) ** 2
													lowerCAmelCase__       =						(self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
													self.assertListEqual(
													    list(attentions[-1].shape[-3:]       )		,    [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]		,    )
													lowerCAmelCase__       =						len(__UpperCAmelCase       )

													# Check attention is always last and order is fine
													lowerCAmelCase__       =						True
													lowerCAmelCase__       =						True
													lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
													model.to(__UpperCAmelCase       )
													model.eval()
													with torch.no_grad():
																			lowerCAmelCase__       =						model(**self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase       )       )

													self.assertEqual(out_len + 1		,    len(__UpperCAmelCase       )       )

													lowerCAmelCase__       =						outputs.attentions

													self.assertEqual(len(__UpperCAmelCase       )		,    __UpperCAmelCase       )
													# verify the first attentions (first block, first layer)
													lowerCAmelCase__       =						(self.model_tester.image_size // 4) ** 2
													lowerCAmelCase__       =						(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]       )		,    [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]		,    )

	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							def check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       ):
													lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
													model.to(__UpperCAmelCase       )
													model.eval()

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

													lowerCAmelCase__       =						outputs.hidden_states

													lowerCAmelCase__       =						self.model_tester.num_encoder_blocks
													self.assertEqual(len(__UpperCAmelCase       )		,    __UpperCAmelCase       )

													# verify the first hidden states (first block)
													self.assertListEqual(
													    list(hidden_states[0].shape[-3:]       )		,    [
													        self.model_tester.hidden_sizes[0],
													        self.model_tester.image_size // 4,
													        self.model_tester.image_size // 4,
													    ]		,    )

							lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													lowerCAmelCase__       =						True
													check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

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

													check_hidden_states_output(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							if not self.model_tester.is_training:
													return

							lowerCAmelCase__      ,							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs_for_common()
							lowerCAmelCase__       =						True

							for model_class in self.all_model_classes:
													if model_class in get_values(__UpperCAmelCase       ):
																			continue

													lowerCAmelCase__       =						model_class(__UpperCAmelCase       )
													model.to(__UpperCAmelCase       )
													model.train()
													lowerCAmelCase__       =						self._prepare_for_class(__UpperCAmelCase		,    __UpperCAmelCase		,    return_labels=__UpperCAmelCase       )
													lowerCAmelCase__       =						model(**__UpperCAmelCase       ).loss
													loss.backward()

	@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."       )
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							pass


	@slow
	def 		UpperCAmelCase  (							self       )->					Union[str, Any]:







							'''simple docstring'''





							for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCAmelCase__       =						SegformerModel.from_pretrained(__UpperCAmelCase       )
													self.assertIsNotNone(__UpperCAmelCase       )





def        _a  (     ) ->   Any:
						"""simple docstring"""







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





@require_torch
class      lowercase__     (			unittest.TestCase	):

	@slow
	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						SegformerImageProcessor(
							    image_scale=(512, 512)		,    keep_ratio=__UpperCAmelCase		,    align=__UpperCAmelCase		,    do_random_crop=__UpperCAmelCase       )
							lowerCAmelCase__       =						SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512"       ).to(
							    __UpperCAmelCase       )

							lowerCAmelCase__       =						prepare_img()
							lowerCAmelCase__       =						image_processor(images=__UpperCAmelCase		,    return_tensors="pt"       )
							lowerCAmelCase__       =						encoded_inputs.pixel_values.to(__UpperCAmelCase       )

							with torch.no_grad():
													lowerCAmelCase__       =						model(__UpperCAmelCase       )

							lowerCAmelCase__       =						torch.Size((1, model.config.num_labels, 128, 128)       )
							self.assertEqual(outputs.logits.shape		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						torch.tensor(
							    [
							        [[-4.6_310, -5.5_232, -6.2_356], [-5.1_921, -6.1_444, -6.5_996], [-5.4_424, -6.2_790, -6.7_574]],
							        [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
							        [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
							    ]       ).to(__UpperCAmelCase       )
							self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,    __UpperCAmelCase		,    atol=1E-4       )       )

	@slow
	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						SegformerImageProcessor(
							    image_scale=(512, 512)		,    keep_ratio=__UpperCAmelCase		,    align=__UpperCAmelCase		,    do_random_crop=__UpperCAmelCase       )
							lowerCAmelCase__       =						SegformerForSemanticSegmentation.from_pretrained(
							    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"       ).to(__UpperCAmelCase       )

							lowerCAmelCase__       =						prepare_img()
							lowerCAmelCase__       =						image_processor(images=__UpperCAmelCase		,    return_tensors="pt"       )
							lowerCAmelCase__       =						encoded_inputs.pixel_values.to(__UpperCAmelCase       )

							with torch.no_grad():
													lowerCAmelCase__       =						model(__UpperCAmelCase       )

							lowerCAmelCase__       =						torch.Size((1, model.config.num_labels, 128, 128)       )
							self.assertEqual(outputs.logits.shape		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						torch.tensor(
							    [
							        [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
							        [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
							        [[-3.6_456, -3.0_209, -1.4_203], [-3.0_797, -3.1_959, -2.0_000], [-1.8_757, -1.9_217, -1.6_997]],
							    ]       ).to(__UpperCAmelCase       )
							self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,    __UpperCAmelCase		,    atol=1E-1       )       )


	@slow
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						SegformerImageProcessor(
							    image_scale=(512, 512)		,    keep_ratio=__UpperCAmelCase		,    align=__UpperCAmelCase		,    do_random_crop=__UpperCAmelCase       )
							lowerCAmelCase__       =						SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512"       ).to(
							    __UpperCAmelCase       )

							lowerCAmelCase__       =						prepare_img()
							lowerCAmelCase__       =						image_processor(images=__UpperCAmelCase		,    return_tensors="pt"       )
							lowerCAmelCase__       =						encoded_inputs.pixel_values.to(__UpperCAmelCase       )

							with torch.no_grad():
													lowerCAmelCase__       =						model(__UpperCAmelCase       )

							lowerCAmelCase__       =						outputs.logits.detach().cpu()

							lowerCAmelCase__       =						image_processor.post_process_semantic_segmentation(outputs=__UpperCAmelCase		,    target_sizes=[(500, 300)]       )
							lowerCAmelCase__       =						torch.Size((500, 300)       )
							self.assertEqual(segmentation[0].shape		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						image_processor.post_process_semantic_segmentation(outputs=__UpperCAmelCase       )
							lowerCAmelCase__       =						torch.Size((128, 128)       )
							self.assertEqual(segmentation[0].shape		,    __UpperCAmelCase       )






