


"""simple docstring"""


import unittest

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


if is_torch_available():
	import torch

	from transformers import AutoModelForImageClassification

if is_vision_available():
	from transformers import AutoImageProcessor
@require_torch
@require_vision
class 						_UpperCamelCase						(     unittest.TestCase	):
							'''simple docstring'''







							@slow
							def     snake_case       (  self			):
										__lowerCAmelCase       					=    AutoImageProcessor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip"			)
										__lowerCAmelCase       					=    AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip"			)
										model.to(__a			)

										from datasets import load_dataset

										__lowerCAmelCase       					=    load_dataset("nielsr/rvlcdip-demo"			)

										__lowerCAmelCase       					=    dataset["train"][0]["image"].convert("RGB"			)

										__lowerCAmelCase       					=    image_processor(__a					,     return_tensors="pt"			).to(__a			)

										# forward pass
										with torch.no_grad():
													__lowerCAmelCase       					=    model(**__a			)
													__lowerCAmelCase       					=    outputs.logits

										__lowerCAmelCase       					=    torch.Size((1, 16)			)
										self.assertEqual(logits.shape					,     __a			)

										__lowerCAmelCase       					=    torch.tensor(
										    [-0.4_1_5_8, -0.4_0_9_2, -0.4_3_4_7]					,     device=__a					,     dtype=torch.float					,     )
										self.assertTrue(torch.allclose(logits[0, :3]					,     __a					,     atol=1e-4			)			)








"""simple docstring"""


import gc
import threading
import time

import psutil
import torch
class 						_UpperCamelCase						:
							'''simple docstring'''







							def __init__(  self			):
										__lowerCAmelCase       					=    psutil.Process()
										__lowerCAmelCase       					=    False

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    -1

										while True:
													__lowerCAmelCase       					=    max(self.process.memory_info().rss					,     self.cpu_memory_peak			)

													# can't sleep or will not catch the peak right (this comment is here on purpose)
													if not self.peak_monitoring:
																break

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    True
										__lowerCAmelCase       					=    threading.Thread(target=self.peak_monitor			)
										__lowerCAmelCase       					=    True
										self.thread.start()



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    False
										self.thread.join()
										return self.cpu_memory_peak
A  :     Any    					=    PeakCPUMemory()


def _lowerCamelCase			(	):

			'''simple docstring'''




			__lowerCAmelCase       					=    {"time": time.time()}

			gc.collect()
			torch.cuda.empty_cache()

			# CPU mem
			__lowerCAmelCase       					=    psutil.Process().memory_info().rss
			cpu_peak_tracker.start()

			# GPU mem
			for i in range(torch.cuda.device_count()   ):
						__lowerCAmelCase       					=    torch.cuda.memory_allocated(_UpperCamelCase   )
			torch.cuda.reset_peak_memory_stats()

			return measures


def _lowerCamelCase			(	_UpperCamelCase   ):

			'''simple docstring'''




			__lowerCAmelCase       					=    {"time": time.time() - start_measures["time"]}

			gc.collect()
			torch.cuda.empty_cache()

			# CPU mem
			__lowerCAmelCase       					=    (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
			__lowerCAmelCase       					=    (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20

			# GPU mem
			for i in range(torch.cuda.device_count()   ):
						__lowerCAmelCase       					=    (torch.cuda.memory_allocated(_UpperCamelCase   ) - start_measures[str(_UpperCamelCase   )]) / 2**20
						__lowerCAmelCase       					=    (torch.cuda.max_memory_allocated(_UpperCamelCase   ) - start_measures[str(_UpperCamelCase   )]) / 2**20

			return measures


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase   ):

			'''simple docstring'''




			print(f"{description}:"   )
			print(f"- Time: {measures['time']:.2f}s"   )
			for i in range(torch.cuda.device_count()   ):
						print(f"- GPU {i} allocated: {measures[str(_UpperCamelCase   )]:.2f}MiB"   )
						__lowerCAmelCase       					=    measures[f"{i}-peak"]
						print(f"- GPU {i} peak: {peak:.2f}MiB"   )
			print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB"   )
			print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB"   )




