ONNXRuntime performance tests

perf_topk.py

import onnxscript as ost
from onnxscript import opset19 as op

import numpy as np

import onnx
import onnxruntime as ort

import time

arr0 = np.array([500], dtype=np.int64)
arr1 = np.array([50], dtype=np.int64)
arr2 = np.array([5], dtype=np.int64)

@ost.script()
def make_topk_2d_axis0(x: ost.FLOAT[1000, 100]) -> (ost.FLOAT[500, 100], ost.INT64[500, 100]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr0)), axis=0)
    return values, indices
@ost.script()
def make_topk_2d_axis0_k5(x: ost.FLOAT[1000, 100]) -> (ost.FLOAT[5, 100], ost.INT64[5, 100]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr2)), axis=0)
    return values, indices

@ost.script()
def make_topk_2d_axis1(x: ost.FLOAT[1000, 100]) -> (ost.FLOAT[1000, 50], ost.INT64[1000, 50]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr1)), axis=1)
    return values, indices

@ost.script()
def make_topk_3d_axis0(x: ost.FLOAT[100, 100, 100]) -> (ost.FLOAT[50, 100, 100], ost.INT64[50, 100, 100]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr1)), axis=0)
    return values, indices
@ost.script()
def make_topk_3d_axis1(x: ost.FLOAT[100, 100, 100]) -> (ost.FLOAT[100, 50, 100], ost.INT64[100, 50, 100]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr1)), axis=1)
    return values, indices
@ost.script()
def make_topk_3d_axis2(x: ost.FLOAT[100, 100, 100]) -> (ost.FLOAT[100, 100, 50], ost.INT64[100, 100, 50]):
    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(arr1)), axis=2)
    return values, indices



models = dict(
    topk_2d_axis0=make_topk_2d_axis0,
    topk_2d_axis0_k5=make_topk_2d_axis0_k5,
    topk_2d_axis1=make_topk_2d_axis1,
    topk_3d_axis0=make_topk_3d_axis0,
    topk_3d_axis1=make_topk_3d_axis1,
    topk_3d_axis2=make_topk_3d_axis2,
)

input_2d_tensor = np.random.randn(1000, 100).astype(np.float32)
input_3d_tensor = np.random.randn(100, 100, 100).astype(np.float32)

def perf_model(k):
    model_proto = models[k].to_model_proto()
    net = ort.InferenceSession(model_proto.SerializeToString(), providers=["CPUExecutionProvider"])

    x = input_2d_tensor if "2d" in k else input_3d_tensor
    print(x.shape)

    for _ in range(10):
        net.run([], {"x": x})

    elapsed_times = []
    for _ in range(100):
        start = time.time()
        net.run([], {"x": x})
        end = time.time()

        elapsed_times.append(1000 * (end - start))

    elapsed_times = sorted(elapsed_times)
    mean = sum(elapsed_times) / len(elapsed_times)
    median = (elapsed_times[49] + elapsed_times[50]) / 2
    minimum = elapsed_times[0]
    print(f"{k}: mean={mean:.2f}, median={median:.2f}, min={minimum:.2f}")

for key in models.keys():
    perf_model(key)