base.py

import torch
import numpy as np
from torchvision import models
from torchvision.datasets import Imagenette
from torch.utils.data import DataLoader
import time
import psutil
import os


model = models.resnet50(weights='IMAGENET1K_V1')
model.eval()

device = torch.device('mps')
model = model.to(device)

dummy_input = torch.from_numpy(np.random.rand(1, 3, 224, 224).astype(np.float32))

N = 100
# ImageNet preprocessing
weights = models.ResNet50_Weights.IMAGENET1K_V1
dataset = Imagenette(root="data", split='val',
                     transform=weights.transforms())
dataset = torch.utils.data.Subset(dataset, range(N))
dataloader = DataLoader(dataset, batch_size=1, shuffle=False)

# warm-up
dummy_input = dummy_input.to(device)
with torch.no_grad():
    for _ in range(10):
        model(dummy_input)

# benchmark
correct = 0
mem_before = psutil.Process(os.getpid()).memory_info().rss
start = time.perf_counter()
for images, labels in dataloader:
    images = images.to(device)
    labels = labels.to(device)
    with torch.no_grad():
        outputs = model(images)
        _, predictions = torch.max(outputs, 1)

    correct += (predictions == labels).sum().item()
end = time.perf_counter()
mem_after = psutil.Process(os.getpid()).memory_info().rss

print(f"PyTorch Accuracy: {correct / N:.2f}")
print(f" Average Inference Time: {(end - start)/N * 1000:.2f} ms")
print(f" Memory Usage: {(mem_after - mem_before) / (1024 * 1024) :.2f} MB")
# 98%, 28 ms, 30~43 MB