onnx_infer_rn50.py

import onnxruntime as ort
import torch
from time import perf_counter


ort_session = ort.InferenceSession("resnet50.onnx", providers=["CPUExecutionProvider"])
print(ort.get_device())
inputs = {}
for inp in ort_session.get_inputs():
    inputs[inp.name] = torch.rand(inp.shape).numpy()

if torch.cuda.is_available():
    # ort_session.set_providers(['CUDAExecutionProvider'])
    print("Switching to GPU")

start = perf_counter()
for _ in range(1000):
    out = ort_session.run(None, inputs)
stop = perf_counter()
print(stop - start)