resnet-50-with-torch-compile.py

# inference using torch.compile()                                                                                                                                                                                                                                                                                                  
# excludes softmax calculation                                                                                                                                                                                                                                                                                                     

import torch
import torchvision.models as models
from torchvision.transforms import functional as F
from PIL import Image
import requests

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


# 1. Load Pretrained ResNet34                                                                                                                                                                                                                                                                                                     \
                                                                                                                                                                                                                                                                                                                                   
model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
model.eval() # put it in eval mode                                                                                                                                                                                                                                                                                                \
                                                                                                                                                                                                                                                                                                                                   

# 2. Compile the model                                                                                                                                                                                                                                                                                                             
# 'reduce-overhead' is great for single image inference                                                                                                                                                                                                                                                                           \
                                                                                                                                                                                                                                                                                                                                   
compiled_model = torch.compile(model, mode="reduce-overhead").to(device="cuda")


# resize and move to cuda                                                                                                                                                                                                                                                                                                          
dim = (224, 224)
print(f"resizing to: {dim}")
img = img.resize(dim)
input_tensor = F.to_tensor(img).unsqueeze(0).to(device="cuda")  # Shape: 1x3x224x224                                                                                                                                                                                                                                               
# processed_input = processor(input_tensor, return_tensors='pt', use_fast=True).to(device="cuda")                                                                                                                                                                                                                                  


print('\nmodel: Resnet34')
torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
    result = model(input_tensor)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
print(f"1st time, includes compile time: {elapsed}")

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

# expected results                                                                                                                                                                                                                                                                                                                 

# ubuntu@147-224-9-53:~$ python r34.py                                                                                                                                                                                                                                                                                             
# image shape:(640, 480)                                                                                                                                                                                                                                                                                                           
# resizing to: (224, 224)                                                                                                                                                                                                                                                                                                          

# model: Resnet34                                                                                                                                                                                                                                                                                                                  
# 1st time, includes compile time: 0.4161966649999158                                                                                                                                                                                                                                                                              
# 2nd time: 0.006660346999979083