Plot onnxruntime profiling

profiling_onnxruntime.py

import torch
import json
import pandas as pd
import matplotlib.pyplot as plt
import os
from pathlib import Path

import onnxruntime

from tqdm import tqdm
import time

import argparse

from optimum.onnxruntime.modeling_ort import ORTModelForImageClassification

parser = argparse.ArgumentParser(description="Baselines")

parser.add_argument("--model", type=str, help="path to model")

parser.add_argument("--threads", type=int, help="path to model")

args, _ = parser.parse_known_args()


model_path = args.model

options = onnxruntime.SessionOptions()
options.intra_op_num_threads = args.threads

options.optimized_model_filepath = os.path.join(os.path.dirname(model_path), Path(model_path).stem + "_optimized.onnx")

options.enable_profiling = True
print(model_path)
print(os.path.join(Path(model_path).stem, Path(model_path).stem + "_optimized.onnx"))

ort_session = onnxruntime.InferenceSession(
    model_path,
    sess_options=options,
    providers=['CPUExecutionProvider']
)

ort_model_eval = ORTModelForImageClassification(ort_session)


inputs = {}
inputs["pixel_values"] = torch.ones(1, 3, 224, 224, dtype=torch.float32)
onnx_inputs = {
    "pixel_values": inputs["pixel_values"].cpu().detach().numpy(),
}

for i in tqdm(range(500)):
    ort_session.run(None, onnx_inputs)

prof = ort_session.end_profiling()

json_path = prof
with open(json_path, "r") as f:
    js = json.load(f)

def process_profiling(js):
    """
    Flattens json returned by onnxruntime profiling.
    :param js: json
    :return: list of dictionaries
    """
    rows = []
    for row in js:
        if 'args' in row and isinstance(row['args'], dict):
            for k, v in row['args'].items():
                row[f'args_{k}'] = v
            del row['args']
        rows.append(row)
    return rows

df = pd.DataFrame(process_profiling(js))

gr_dur = df[['dur', "args_op_name"]].groupby(
    "args_op_name").sum().sort_values('dur')
gr_n = df[['dur', "args_op_name"]].groupby(
    "args_op_name").count().sort_values('dur')
gr_n = gr_n.loc[gr_dur.index, :]

fig, ax = plt.subplots(1, 2, figsize=(8, 4))
gr_dur.plot.barh(ax=ax[0])
gr_n.plot.barh(ax=ax[1])
ax[0].set_title("duration")
ax[1].set_title("n occurences")

plt.savefig(os.path.join(os.path.dirname(model_path), f"{time.strftime('%Y%m%d-h%Hm%Ms%S')}_profile_" + Path(model_path).stem + ".png"))