onnx_optimized.py

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


opts = ort.SessionOptions()
opts.intra_op_num_threads = 4
opts.execution_mode = ort.ExecutionMode.ORT_PARALLEL
session = ort.InferenceSession("models/resnet50_fp32.onnx",
                               sess_options=opts,
                               providers=["CoreMLExecutionProvider"])
onnx_input_name = session.get_inputs()[0].name

dummy_input = 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
for _ in range(10):
    session.run(None, {onnx_input_name: dummy_input})

# benchmark
correct = 0
mem_before = psutil.Process(os.getpid()).memory_info().rss
start = time.time()
for images, labels in dataloader:
    outputs = session.run(None, {onnx_input_name: np.array(images)})
    prediction = np.argmax(outputs[0])
    correct += (prediction == labels.item())
end = time.time()
mem_after = psutil.Process(os.getpid()).memory_info().rss

print(f"ONNX 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")
# precision f16: 98%, 78-80 ms, 35-50MB (includes converting data to flaot16)
# precision f32: 98%, 11-13 ms, 10-14 MB