inference-speed-vit-detr-A10.py

 Dependencies for test.py (ViT + DETR benchmarks)                                                                                                                                                                                                             
# Tested with Python 3.10 on Ubuntu / NVIDIA A10 (CUDA 12.8)                                                                                                                                                                                                   

# torch==2.7.0                                                                                                                                                                                                                                                 
# transformers==5.8.0                                                                                                                                                                                                                                          
# Pillow >= 9.1.0 is required by transformers (PIL.Image.Resampling)                                                                                                                                                                                           
# Pillow>=10.0.0                                                                                                                                                                                                                                               
# timm is required by transformers backbones such as facebook/detr-resnet-50                                                                                                                                                                                   
# timm>=1.0.0                                                                                                                                                                                                                                                  
# huggingface_hub>=1.0.0                                                                                                                                                                                                                                       
# numpy>=1.21                                                                                                                                                                                                                                                  
# requests>=2.25                                                                                                                                                                                                                                               

# rough benchmarks of compile time and run time for ViT and detr-R-50                                                                                                                                                                                         \
                                                                                                                                                                                                                                                               
import torch
from PIL import Image
import requests
import numpy as np
from transformers import AutoImageProcessor, AutoModelForImageClassification

print(f"cuda version: {torch.version.cuda}")
if torch.cuda.is_available():
    # Print the name of the first GPU (device 0)                                                                                                                                                                                                              \
                                                                                                                                                                                                                                                               
    print(f"GPU Flavor: {torch.cuda.get_device_name(0)}")
else:
    print("CUDA is not available. Running on CPU.")

# torch._logging.set_logs(graph_code=True)                                                                                                                                                                                                                    \
                                                                                                                                                                                                                                                               

import torch, time
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
print(f"image shape:{image.size}")

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True)
model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224").to("cuda")
model = torch.compile(model)

# move to cuda                                                                                                                                                                                                                                                \
                                                                                                                                                                                                                                                               
processed_input = processor(image, return_tensors='pt').to(device="cuda")

print('\nmodel: ViT')
torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
    result = model(**processed_input)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"1st time: {elapsed}")

# lazy compile do it again                                                                                                                                                                                                                                    \
                                                                                                                                                                                                                                                               
torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
    result = model(**processed_input)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"2nd time: {elapsed}")


from transformers import AutoImageProcessor, AutoModelForObjectDetection
processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50").to("cuda")
model = torch.compile(model)

# inputs = processor(text=texts, images=image, return_tensors="pt").to("cuda")                                                                                                                                                                                \
                                                                                                                                                                                                                                                               
inputs = processor(images=image, return_tensors="pt").to("cuda")

print('\nmodel: detr-resnet-50')
torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
    _ = model(**inputs)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"1st time: {elapsed}")

torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
    _ = model(**inputs)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"2nd time: {elapsed}")