[Python] Abnormal cuda and cpu memory usage with huggingface transformers ViT-msn model

There are basically two parts to the question, the first part is that while inferring on the model and storing minimal results, the python app takes up around 3GB of memory on my system ( i have hardcoded device='cpu'), this is unexpected as the model is a vit-msn-base model with 12 layers only and is around 160mb in size, my batch size is 1.. Also, even on using gc.collect(), it seems like the memory usage just keeps increasing

the second part is that the memory usage spikes to around 10GB when i use device='cuda', this was really unexpected, on some googling, it seems like pytorch's memory handling results in extensive fragmentation, but i am not sure on how to handle that

here are the relevant code snippets -

model_name = 'facebook/vit-msn-base-4'
layers = 1


cka_score_list = []

tokenizer = AutoImageProcessor.from_pretrained(
model_name,
use_fast=True, # Use the fast tokenizer implementation
trust_remote_code=True, # Trust remote code (required for some models)
add_bos_token=False, # Do not add beginning-of-sequence token
add_eos_token=False, # Do not add end-of-sequence token
padding_side="left" # Pad sequences on the left side
)

# Load the pre-trained causal language model with appropriate settings
model = AutoModel.from_pretrained(
model_name,
trust_remote_code=True, # Trust remote code (required for some models)
device_map="auto", # Automatically map layers to available devices
torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32, # Use bfloat16 if supported
)
print(dir(model))
model.eval()
model = model.to(device)
cur_score = eval_meth(classA, classB, num_images)

cka_score_list.append(cur_score)


def latent_embeddings_for_class(classA, num_images):
cur_csv = os.path.join(val_csv , (classA+ '.csv'))
df = pd.read_csv(cur_csv, header = None)
embeddings = []
for i in range(num_images):
image_name = df.iloc[random.randint(1, 50), 0]
image_path = os.path.join(val_images, image_name)
image =Image.open(image_path)

inputs = tokenizer(image, return_tensors="pt").to(device).to(torch.bfloat16)
print(inputs['pixel_values'].shape)
snapshot = tracemalloc.take_snapshot()

# Print top memory-consuming lines
top_stats = snapshot.statistics('lineno')
for stat in top_stats[:3]:
print(stat)

output = model(**inputs)
embeddings.append(output.last_hidden_state[0,1])

del inputs
del image
del output
clear_memory()

print(embeddings[0].shape)

embeddings = torch.stack(embeddings)
return embeddings



def eval_meth(classA, classB, num_images ):

X= latent_embeddings_for_class( classA, num_images)
Y= latent_embeddings_for_class( classB, num_images)

cka_score = cka(X,Y)
clear_memory()
return cka_score


also, on CPU, this is the output -

torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\selectors.py:315: size=288 KiB, count=6, average=48.0 KiB
torch.Size([768])
torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\tracemalloc.py:505: size=537 KiB, count=9815, average=56 B
torch.Size([768])
torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\tracemalloc.py:505: size=533 KiB, count=9732, average=56 B
torch.Size([768])


on GPU it shows -

CUDA out of memory. Tried to allocate 452.00 MiB. GPU 0 has a total capacity of 4.00 GiB of which 0 bytes is free. Of the allocated memory 10.17 GiB is allocated by PyTorch, and 144.56 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


this is weird as this does not reflect the memory usage shown in task manager

Continue reading...