vit.py

import torch
from time import time

import numpy as np
import os
import sys
from torch.profiler import profile, record_function, ProfilerActivity
times=[]

from transformers import AutoImageProcessor, ViTForImageClassification
from datasets import load_dataset

dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

inputs = image_processor(image, return_tensors="pt")

## to enable torch inductor uncomment below line
model = torch.compile(model)

with torch.set_grad_enabled(False):
    for _ in range(50):
        model(**inputs) #Warmup
    print("Warmup over")
    with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
        with record_function("model_inference"):
            for _ in range(100):
                model(**inputs)

print(prof.key_averages(group_by_input_shape=True).table(sort_by="self_cpu_time_total"))