201: Evaluate DORT

It compares DORT to eager mode and onnxrt backend.

To run the script:

python _doc/examples/plot_torch_dort --help

Some helpers

import warnings

try:
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        import onnxruntime

        has_cuda = "CUDAExecutionProvider" in onnxruntime.get_available_providers()
except ImportError:
    print("onnxruntime not available.")
    import sys

    sys.exit(0)

import torch._dynamo
import contextlib
import itertools
import os
import gc
import platform

# import pickle
import pprint
import multiprocessing
import time
import cProfile
import pstats
import io
import logging
from pstats import SortKey

import numpy as np
import matplotlib.pyplot as plt
import pandas
import onnx
from onnx_array_api.profiling import profile2graph
import torch
from torch import nn
import torch.nn.functional as F
import experimental_experiment
from experimental_experiment.plotting.memory import memory_peak_plot
from experimental_experiment.ext_test_case import measure_time, get_figure
from experimental_experiment.args import get_parsed_args
from experimental_experiment.memory_peak import start_spying_on
from experimental_experiment.torch_models.training_helper import make_aot_ort
from tqdm import tqdm

has_cuda = has_cuda and torch.cuda.is_available()
logging.disable(logging.ERROR)


def system_info():
    obs = {}
    obs["processor"] = platform.processor()
    obs["cores"] = multiprocessing.cpu_count()
    try:
        obs["cuda"] = 1 if torch.cuda.is_available() else 0
        obs["cuda_count"] = torch.cuda.device_count()
        obs["cuda_name"] = torch.cuda.get_device_name()
        obs["cuda_capa"] = torch.cuda.get_device_capability()
    except (RuntimeError, AssertionError):
        # no cuda
        pass
    return obs


pprint.pprint(system_info())
{'cores': 20,
 'cuda': 1,
 'cuda_capa': (8, 9),
 'cuda_count': 1,
 'cuda_name': 'NVIDIA GeForce RTX 4060 Laptop GPU',
 'processor': 'x86_64'}

Scripts arguments

script_args = get_parsed_args(
    "plot_torch_dort",
    description=__doc__,
    scenarios={
        "small": "small model to test",
        "middle": "55Mb model",
        "large": "1Gb model",
    },
    warmup=5,
    repeat=5,
    repeat1=(1, "repeat for the first iteration"),
    maxtime=(
        2,
        "maximum time to run a model to measure the computation time, "
        "it is 0.1 when scenario is small",
    ),
    expose="scenarios,repeat,repeat1,warmup",
)

if script_args.scenario in (None, "small"):
    script_args.maxtime = 0.1
print(f"scenario={script_args.scenario or 'small'}")
print(f"warmup={script_args.warmup}")
print(f"repeat={script_args.repeat}")
print(f"repeat1={script_args.repeat1}")
print(f"maxtime={script_args.maxtime}")
scenario=small
warmup=5
repeat=5
repeat1=1
maxtime=0.1

The model

A simple model to convert.

class MyModelClass(nn.Module):
    def __init__(self, scenario=script_args.scenario):
        super().__init__()
        if scenario == "middle":
            self.large = False
            self.conv1 = nn.Conv2d(1, 32, 5)
            # self.conv2 = nn.Conv2d(128, 16, 5)
            self.fc1 = nn.Linear(30752, 1024)
            self.fcs = []
            self.fc2 = nn.Linear(1024, 128)
            self.fc3 = nn.Linear(128, 10)
        elif scenario in (None, "small"):
            self.large = False
            self.conv1 = nn.Conv2d(1, 16, 5)
            # self.conv2 = nn.Conv2d(16, 16, 5)
            self.fc1 = nn.Linear(144, 512)
            self.fcs = []
            self.fc2 = nn.Linear(512, 128)
            self.fc3 = nn.Linear(128, 10)
        elif scenario in (None, "large"):
            self.large = True
            self.conv1 = nn.Conv2d(1, 32, 5)
            # self.conv2 = nn.Conv2d(128, 16, 5)
            self.fc1 = nn.Linear(30752, 4096)
            # torch script does not support loops.
            self.fca = nn.Linear(4096, 4096)
            self.fcb = nn.Linear(4096, 4096)
            self.fcc = nn.Linear(4096, 4096)
            self.fcd = nn.Linear(4096, 4096)
            self.fce = nn.Linear(4096, 4096)
            self.fcf = nn.Linear(4096, 4096)
            self.fcg = nn.Linear(4096, 4096)
            self.fch = nn.Linear(4096, 4096)
            self.fci = nn.Linear(4096, 4096)
            # end of the unfolded loop.
            self.fc2 = nn.Linear(4096, 128)
            self.fc3 = nn.Linear(128, 10)
        else:
            raise ValueError(f"Unsupported scenario={scenario!r}.")

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (4, 4))
        # x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        if self.large:
            # loop
            x = F.relu(self.fca(x))
            x = F.relu(self.fcb(x))
            x = F.relu(self.fcc(x))
            x = F.relu(self.fcd(x))
            x = F.relu(self.fce(x))
            x = F.relu(self.fcf(x))
            x = F.relu(self.fcg(x))
            x = F.relu(self.fch(x))
            x = F.relu(self.fci(x))
            # end of the loop
        x = F.relu(self.fc2(x))
        y = self.fc3(x)
        return y


def create_model_and_input(scenario=script_args.scenario):
    if scenario == "middle":
        shape = [1, 1, 128, 128]
    elif scenario in (None, "small"):
        shape = [1, 1, 16, 16]
    elif scenario == "large":
        shape = [1, 1, 128, 128]
    else:
        raise ValueError(f"Unsupported scenario={scenario!r}.")
    input_tensor = torch.rand(*shape).to(torch.float32)
    model = MyModelClass(scenario=scenario)
    assert model(input_tensor) is not None
    return model, input_tensor


def torch_model_size(model):
    size_model = 0
    for param in model.parameters():
        size = param.numel() * torch.finfo(param.data.dtype).bits / 8
        size_model += size
    return size_model


model, input_tensor = create_model_and_input()
model_size = torch_model_size(model)
print(f"model size={model_size / 2 ** 20} Mb")
model size=0.5401992797851562 Mb

Backends

def get_torch_eager(model, *args):
    def my_compiler(gm, example_inputs):
        return gm.forward

    with contextlib.redirect_stdout(io.StringIO()):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            optimized_mod = torch.compile(model, fullgraph=True, backend=my_compiler)
            optimized_mod(*args)
            return optimized_mod


def get_torch_default(model, *args):
    with contextlib.redirect_stdout(io.StringIO()):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            optimized_mod = torch.compile(model, fullgraph=True, mode="reduce-overhead")
            optimized_mod(*args)
            return optimized_mod


def get_torch_dort(model, *args):
    with contextlib.redirect_stdout(io.StringIO()):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            local_aot_ort, _ = make_aot_ort(dynamic=True, rewrite=True)
            optimized_mod = torch.compile(model, backend=local_aot_ort, fullgraph=True)
            optimized_mod(*args)
            return optimized_mod


def get_torch_opti(model, *args):
    with contextlib.redirect_stdout(io.StringIO()):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            local_aot_ort, _ = make_aot_ort(dynamic=True, rewrite=True)
            optimized_mod = torch.compile(model, backend=local_aot_ort, fullgraph=True)
            optimized_mod(*args)
            return optimized_mod

Let’s check they are working.

export_functions = [
    get_torch_eager,
    get_torch_default,
    get_torch_dort,
    # get_torch_opti,
]

exporters = {f.__name__.replace("get_", ""): f for f in export_functions}

supported_exporters = {}
for k, v in exporters.items():
    print(f"run function {k}")
    filename = f"plot_torch_dort_{k}.onnx"
    torch._dynamo.reset()
    model, input_tensor = create_model_and_input()
    try:
        v(model, input_tensor)
    except Exception as e:
        print(f"skipped due to {str(e)[:1000]}")
        continue
    supported_exporters[k] = v
    del model
    gc.collect()
    time.sleep(1)
run function torch_eager
run function torch_default
run function torch_dort

Compile and Memory

def flatten(ps):
    obs = ps["cpu"].to_dict(unit=2**20)
    if "gpus" in ps:
        for i, g in enumerate(ps["gpus"]):
            for k, v in g.to_dict(unit=2**20).items():
                obs[f"gpu{i}_{k}"] = v
    return obs


data = []

for k, v in supported_exporters.items():
    print(f"run compile for memory {k} on cpu")
    filename = f"plot_torch_dort_{k}.onnx"
    if has_cuda:
        torch.cuda.set_device(0)
    torch._dynamo.reset()
    # CPU
    model, input_tensor = create_model_and_input()
    stat = start_spying_on(cuda=1 if has_cuda else 0)
    v(model, input_tensor)
    obs = flatten(stat.stop())
    print("done.")
    obs.update(dict(export=k, p="cpu"))
    data.append(obs)
    del model
    gc.collect()
    time.sleep(1)

    if not has_cuda:
        continue
    if k in {"torch_default"}:
        print(f"skip compile for memory {k} on cuda")
        continue
    torch._dynamo.reset()
    # CUDA
    model, input_tensor = create_model_and_input()
    model = model.cuda()
    input_tensor = input_tensor.cuda()
    print(f"run compile for memory {k} on cuda")
    stat = start_spying_on(cuda=1 if has_cuda else 0)
    v(model, input_tensor)
    obs = flatten(stat.stop())
    print("done.")
    obs.update(dict(export=k, p="cuda"))
    data.append(obs)
    del model
    gc.collect()
    time.sleep(1)
run compile for memory torch_eager on cpu
done.
run compile for memory torch_eager on cuda
done.
run compile for memory torch_default on cpu
done.
skip compile for memory torch_default on cuda
run compile for memory torch_dort on cpu
done.
run compile for memory torch_dort on cuda
done.

The result.

df1 = pandas.DataFrame(data)
df1.to_csv("plot_torch_dort_1_memory.csv", index=False)
df1.to_excel("plot_torch_dort_1_memory.xlsx", index=False)
print(df1)

for p in ["cpu", "cuda"]:
    if not has_cuda and p == "cuda":
        continue
    ax = memory_peak_plot(
        df1[df1["p"] == p],
        key=("export",),
        bars=[model_size * i / 2**20 for i in range(1, 5)],
        suptitle=f"Memory Consumption of the Compilation on {p}\n"
        f"model size={model_size / 2**20:1.0f} Mb",
    )
    get_figure(ax).savefig(f"plot_torch_dort_1_memory_{p}.png")
  • Memory Consumption of the Compilation on cpu model size=1 Mb, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)
  • Memory Consumption of the Compilation on cuda model size=1 Mb, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)
          peak         mean   n        begin          end   gpu0_peak   gpu0_mean  gpu0_n  gpu0_begin    gpu0_end         export     p
0  3323.484375  3323.153646  12  3323.484375  3321.503906  876.617188  876.617188      12  876.617188  876.617188    torch_eager   cpu
1  3321.457031  3321.457031  11  3321.457031  3321.457031  896.617188  878.435369      11  876.617188  896.617188    torch_eager  cuda
2  3321.445312  3321.445186  31  3321.445312  3321.441406  896.617188  896.617188      31  896.617188  896.617188  torch_default   cpu
3  3325.480469  3323.484855  57  3323.449219  3325.480469  896.617188  896.617188      57  896.617188  896.617188     torch_dort   cpu
4  3325.511719  3325.477202  55  3325.476562  3325.511719  904.617188  896.762642      55  896.617188  904.617188     torch_dort  cuda

dort first iteration speed

data = []

for k, v in supported_exporters.items():
    print(f"run dort cpu {k}: {script_args.repeat1}")
    times = []
    for _ in range(int(script_args.repeat1)):
        model, input_tensor = create_model_and_input()
        torch._dynamo.reset()
        begin = time.perf_counter()
        v(model, input_tensor)
        duration = time.perf_counter() - begin
        times.append(duration)
        del model
        gc.collect()
        time.sleep(1)

    print(f"done: {times[-1]}")
    data.append(
        dict(
            export=k,
            time=np.mean(times),
            min=min(times),
            max=max(times),
            first=times[0],
            last=times[-1],
            std=np.std(times),
            p="cpu",
        )
    )

    if not has_cuda:
        continue
    if k in {"torch_dort", "torch_default"}:
        print(f"skip dort cuda {k}: {script_args.repeat1}")
        continue
    print(f"run dort cuda {k}: {script_args.repeat1}")
    times = []
    for _ in range(int(script_args.repeat1)):
        model, input_tensor = create_model_and_input()
        model = model.cuda()
        input_tensor = input_tensor.cuda()
        torch._dynamo.reset()
        begin = time.perf_counter()
        v(model, input_tensor)
        duration = time.perf_counter() - begin
        times.append(duration)
        del model
        gc.collect()
        time.sleep(1)

    print(f"done: {times[-1]}")
    data.append(
        dict(
            export=k,
            time=np.mean(times),
            min=min(times),
            max=max(times),
            first=times[0],
            last=times[-1],
            std=np.std(times),
            p="cuda",
        )
    )
run dort cpu torch_eager: 1
done: 0.08234240999809117
run dort cuda torch_eager: 1
done: 0.04043209900191869
run dort cpu torch_default: 1
done: 0.1819591470011801
skip dort cuda torch_default: 1
run dort cpu torch_dort: 1
done: 0.4098883000006026
skip dort cuda torch_dort: 1

The result.

df1 = pandas.DataFrame(data)
df1.to_csv("plot_torch_dort_1_time.csv", index=False)
df1.to_excel("plot_torch_dort_1_time.xlsx", index=False)
print(df1)

fig, ax = plt.subplots(1, 1)
dfi = df1[["export", "p", "time", "std"]].set_index(["export", "p"])
dfi["time"].plot.bar(ax=ax, title="Compilation time", yerr=dfi["std"], rot=30)
fig.tight_layout()
fig.savefig("plot_torch_dort_1_time.png")
Compilation time
          export      time       min       max     first      last  std     p
0    torch_eager  0.082342  0.082342  0.082342  0.082342  0.082342  0.0   cpu
1    torch_eager  0.040432  0.040432  0.040432  0.040432  0.040432  0.0  cuda
2  torch_default  0.181959  0.181959  0.181959  0.181959  0.181959  0.0   cpu
3     torch_dort  0.409888  0.409888  0.409888  0.409888  0.409888  0.0   cpu

Compilation Profiling

def clean_text(text):
    pathes = [
        os.path.abspath(os.path.normpath(os.path.join(os.path.dirname(torch.__file__), ".."))),
        os.path.abspath(os.path.normpath(os.path.join(os.path.dirname(onnx.__file__), ".."))),
        os.path.abspath(
            os.path.normpath(
                os.path.join(os.path.dirname(experimental_experiment.__file__), "..")
            )
        ),
    ]
    for p in pathes:
        text = text.replace(p, "")
    text = text.replace("experimental_experiment", "experimental_experiment".upper())
    return text


def profile_function(name, export_function, with_args=True, verbose=False, suffix="export"):
    if verbose:
        print(f"profile {name}: {export_function}")
    if with_args:
        model, input_tensor = create_model_and_input()
        pr = cProfile.Profile()
        pr.enable()
        for _ in range(int(script_args.repeat1)):
            export_function(model, input_tensor)
        pr.disable()
    else:
        pr = cProfile.Profile()
        pr.enable()
        for _ in range(int(script_args.repeat1)):
            export_function()
        pr.disable()
    s = io.StringIO()
    sortby = SortKey.CUMULATIVE
    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
    ps.print_stats()
    # with open(f"plot_torch_dort_profile_{name}_{suffix}.pickle", "wb") as f:
    #     pickle.dump(ps, f)

    raw = s.getvalue()
    text = "\n".join(raw.split("\n")[:200])
    if verbose:
        print(text)
    with open(f"plot_torch_dort_profile_{name}_{suffix}.txt", "w") as f:
        f.write(raw)

    root, nodes = profile2graph(ps, clean_text=clean_text)
    text = root.to_text()
    with open(f"plot_torch_dort_profile_{name}_{suffix}_h.txt", "w") as f:
        f.write(text)
    if verbose:
        print("done.")


model, input_tensor = create_model_and_input()


def function_to_profile(model=model, input_tensor=input_tensor):
    return get_torch_dort(model, input_tensor)


profile_function("dort", function_to_profile, verbose=True, suffix="1")
profile dort: <function function_to_profile at 0x7f0082fd4310>
         1093346 function calls (1065254 primitive calls) in 0.755 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.793    0.793 /home/xadupre/github/experimental-experiment/_doc/examples/plot_torch_dort_201.py:506(function_to_profile)
        1    0.000    0.000    0.793    0.793 /home/xadupre/github/experimental-experiment/_doc/examples/plot_torch_dort_201.py:240(get_torch_dort)
      4/1    0.000    0.000    0.413    0.413 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/nn/modules/module.py:1732(_wrapped_call_impl)
      4/1    0.000    0.000    0.413    0.413 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/nn/modules/module.py:1740(_call_impl)
        1    0.000    0.000    0.413    0.413 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:523(_fn)
        1    0.000    0.000    0.378    0.378 /home/xadupre/github/experimental-experiment/experimental_experiment/torch_models/training_helper.py:7(make_aot_ort)
        1    0.000    0.000    0.378    0.378 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/onnxruntime.py:763(__init__)
        1    0.000    0.000    0.325    0.325 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:308(__init__)
        1    0.000    0.000    0.321    0.321 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1331(__call__)
        1    0.000    0.000    0.320    0.320 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:449(__call__)
        1    0.000    0.000    0.320    0.320 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:597(_compile)
        1    0.000    0.000    0.319    0.319 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:689(compile_inner)
        1    0.000    0.000    0.319    0.319 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_utils_internal.py:89(wrapper_function)
        1    0.000    0.000    0.319    0.319 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:709(_compile_inner)
        1    0.000    0.000    0.301    0.301 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py:1329(transform_code_object)
        1    0.000    0.000    0.299    0.299 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:204(_fn)
        1    0.000    0.000    0.299    0.299 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:632(transform)
        1    0.000    0.000    0.297    0.297 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:2907(run)
      6/1    0.000    0.000    0.297    0.297 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1110(run)
   100/44    0.001    0.000    0.297    0.007 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:998(step)
        1    0.000    0.000    0.264    0.264 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:95(__init__)
        1    0.001    0.001    0.264    0.264 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:123(_initiate_registry_from_torchlib)
        1    0.005    0.005    0.261    0.261 /home/xadupre/github/onnxscript/onnxscript/_framework_apis/torch_2_5.py:107(get_torchlib_ops)
      184    0.001    0.000    0.255    0.001 /home/xadupre/github/onnxscript/onnxscript/values.py:640(function_ir)
        1    0.000    0.000    0.237    0.237 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:3098(RETURN_VALUE)
        1    0.000    0.000    0.237    0.237 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:3070(_return)
        1    0.000    0.000    0.237    0.237 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/output_graph.py:989(compile_subgraph)
        1    0.000    0.000    0.236    0.236 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/output_graph.py:1324(compile_and_call_fx_graph)
        1    0.000    0.000    0.231    0.231 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/output_graph.py:1444(call_user_compiler)
        1    0.000    0.000    0.231    0.231 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/output_graph.py:1450(_call_user_compiler)
      2/1    0.000    0.000    0.230    0.230 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/repro/after_dynamo.py:73(__call__)
        1    0.000    0.000    0.230    0.230 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/__init__.py:2318(__call__)
        1    0.000    0.000    0.230    0.230 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/onnxruntime.py:1153(__call__)
        1    0.000    0.000    0.230    0.230 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/backends/common.py:23(__call__)
        1    0.000    0.000    0.230    0.230 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:940(aot_module_simplified)
        1    0.000    0.000    0.225    0.225 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1061(dispatch_and_compile)
        1    0.000    0.000    0.225    0.225 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:516(create_aot_dispatcher_function)
        1    0.000    0.000    0.225    0.225 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:529(_create_aot_dispatcher_function)
      3/2    0.000    0.000    0.196    0.098 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:715(_fn)
        1    0.000    0.000    0.193    0.193 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:337(aot_dispatch_autograd)
      184    0.001    0.000    0.168    0.001 /home/xadupre/github/onnxscript/onnxscript/_internal/ast_utils.py:16(get_src_and_ast)
      185    0.000    0.000    0.140    0.001 /usr/lib/python3.10/inspect.py:1133(getsource)
      185    0.004    0.000    0.139    0.001 /usr/lib/python3.10/inspect.py:1112(getsourcelines)
      184    0.020    0.000    0.124    0.001 /usr/lib/python3.10/inspect.py:1101(getblock)
        1    0.000    0.000    0.112    0.112 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py:234(aot_dispatch_autograd_graph)
        2    0.025    0.012    0.108    0.054 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/decomposition_table.py:14(_create_onnx_supports_op_overload_table)
      155    0.005    0.000    0.106    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/functional_tensor.py:372(__torch_dispatch__)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py:46(_create_graph)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:2170(wrapped)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:2108(trace)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:1999(_trace_inner)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_compile.py:22(inner)
        1    0.000    0.000    0.104    0.104 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:1131(dispatch_trace)
        1    0.000    0.000    0.102    0.102 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:711(trace)
        1    0.000    0.000    0.099    0.099 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:698(flatten_fn)
        1    0.000    0.000    0.099    0.099 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:1181(wrapped)
        1    0.000    0.000    0.096    0.096 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:663(inner_fn)
        1    0.000    0.000    0.096    0.096 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:643(joint_helper)
        1    0.000    0.000    0.096    0.096 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:396(_functionalized_f_helper)
    26710    0.054    0.000    0.094    0.000 /usr/lib/python3.10/tokenize.py:431(_tokenize)
        1    0.000    0.000    0.093    0.093 /home/xadupre/github/experimental-experiment/_doc/examples/plot_torch_dort_201.py:165(forward)
        1    0.000    0.000    0.093    0.093 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py:1113(forward)
        1    0.000    0.000    0.093    0.093 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:284(runtime_wrapper)
      2/1    0.000    0.000    0.093    0.093 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py:116(call_func_at_runtime_with_args)
      2/1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py:99(g)
        1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/autograd/function.py:559(apply)
        1    0.000    0.000    0.092    0.092 {built-in method apply}
        1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:1520(forward)
        1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:474(wrapper)
        1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:659(inner_fn)
        1    0.000    0.000    0.092    0.092 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py:115(_lazy_forward)
  815/495    0.001    0.000    0.091    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_stats.py:16(wrapper)
      2/1    0.000    0.000    0.091    0.091 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph_module.py:821(call_wrapped)
        1    0.000    0.000    0.091    0.091 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph_module.py:382(__call__)
        1    0.000    0.000    0.091    0.091 <eval_with_key>.320:4(forward)
        1    0.000    0.000    0.091    0.091 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/onnxruntime.py:884(_ort_acclerated_call)
      429    0.001    0.000    0.086    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:1230(__torch_function__)
        1    0.000    0.000    0.085    0.085 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:276(inner_fn_with_anomaly)
        1    0.000    0.000    0.085    0.085 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:193(inner_fn)
      184    0.000    0.000    0.078    0.000 /home/xadupre/github/onnxscript/onnxscript/converter.py:1463(translate_function_signature)
      184    0.006    0.000    0.077    0.000 /home/xadupre/github/onnxscript/onnxscript/converter.py:1378(_translate_function_signature_common)
        5    0.000    0.000    0.067    0.013 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/overrides.py:1668(handle_torch_function)
  297/276    0.001    0.000    0.062    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:1328(__torch_dispatch__)
        1    0.001    0.001    0.060    0.060 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/decomposition_table.py:73(create_onnx_friendly_decomposition_table)
153601/149275    0.022    0.000    0.057    0.000 {built-in method builtins.isinstance}
        1    0.000    0.000    0.055    0.055 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/partitioners.py:1779(min_cut_rematerialization_partition)
    69/54    0.002    0.000    0.055    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:761(proxy_call)
     14/9    0.000    0.000    0.053    0.006 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:627(wrapper)
     14/9    0.000    0.000    0.053    0.006 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1728(CALL_FUNCTION)
      2/1    0.000    0.000    0.053    0.053 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/autograd/__init__.py:358(grad)
     14/9    0.000    0.000    0.053    0.006 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:941(call_function)
        1    0.000    0.000    0.053    0.053 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/autograd/graph.py:816(_engine_run_backward)
        1    0.002    0.002    0.053    0.053 {method 'run_backward' of 'torch._C._EngineBase' objects}
        3    0.000    0.000    0.053    0.018 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/interpreter.py:117(run)
  516/511    0.001    0.000    0.051    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1242(__torch_dispatch__)
  516/511    0.003    0.000    0.050    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1768(dispatch)
       65    0.000    0.000    0.049    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/interpreter.py:210(run_node)
    16562    0.007    0.000    0.048    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:210(is_registered_op)
7575/1820    0.009    0.000    0.047    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:131(is_value_type)
       35    0.000    0.000    0.047    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/interpreter.py:288(call_function)
     17/6    0.000    0.000    0.046    0.008 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py:166(realize_and_forward)
        4    0.000    0.000    0.046    0.012 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py:850(call_function)
      150    0.001    0.000    0.046    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1326(_cached_dispatch_impl)
        2    0.000    0.000    0.046    0.023 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:848(functional_call)
       38    0.000    0.000    0.044    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py:6471(run_node)
      5/4    0.000    0.000    0.043    0.011 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:325(call_function)
      5/4    0.000    0.000    0.043    0.011 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:119(call_function)
      5/4    0.000    0.000    0.043    0.011 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:963(inline_user_function_return)
      5/4    0.000    0.000    0.043    0.011 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:3120(inline_call)
      5/4    0.000    0.000    0.043    0.011 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:3157(inline_call_)
    16577    0.010    0.000    0.040    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/_exporter_legacy.py:188(get_op_functions)
        1    0.000    0.000    0.038    0.038 /home/xadupre/github/onnxscript/onnxscript/optimizer/__init__.py:15(optimize)
        1    0.000    0.000    0.038    0.038 /home/xadupre/github/onnxscript/onnxscript/optimizer/_legacy/_optimizer.py:24(optimize)
  533/436    0.000    0.000    0.034    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_ops.py:722(__call__)
  328/144    0.001    0.000    0.034    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_pytree.py:923(tree_map)
     1569    0.032    0.000    0.032    0.000 {built-in method builtins.compile}
        1    0.000    0.000    0.032    0.032 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py:109(inner_fn)
 1500/221    0.003    0.000    0.031    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_pytree.py:801(unflatten)
     1024    0.001    0.000    0.031    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:172(is_valid_type)
     87/3    0.001    0.000    0.030    0.010 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/diagnostics/infra/decorator.py:66(wrapper)
       28    0.000    0.000    0.030    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1869(LOAD_ATTR)
       28    0.000    0.000    0.030    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1862(_load_attr)
    31/29    0.000    0.000    0.029    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:980(call_function)
       75    0.000    0.000    0.029    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/base.py:366(build)
       29    0.000    0.000    0.028    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:843(builtin_dispatch)
       28    0.000    0.000    0.028    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:763(call_self_handler)
       28    0.001    0.000    0.028    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:1615(call_getattr)
       75    0.000    0.000    0.028    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:371(__call__)
      486    0.002    0.000    0.027    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph.py:1104(create_node)
    25618    0.026    0.000    0.026    0.000 {method 'match' of 're.Pattern' objects}
    51/35    0.000    0.000    0.026    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py:1073(var_getattr)
       40    0.001    0.000    0.026    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:536(_wrap)
  897/801    0.002    0.000    0.026    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:150(<listcomp>)
        4    0.001    0.000    0.025    0.006 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/partitioners.py:153(_extract_graph_with_inputs_outputs)
    35/19    0.001    0.000    0.025    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py:993(var_getattr)
      141    0.000    0.000    0.025    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py:61(realize)
       36    0.000    0.000    0.025    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py:20(realize)
1281/1131    0.001    0.000    0.024    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/base.py:102(__instancecheck__)
      314    0.001    0.000    0.024    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_pytree.py:1130(tree_map_only)
        1    0.000    0.000    0.024    0.024 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py:171(inner)
      327    0.001    0.000    0.023    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph.py:1493(node_copy)
       18    0.000    0.000    0.023    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:2079(wrap_fx_proxy)
       18    0.001    0.000    0.023    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:2141(wrap_fx_proxy_cls)
      184    0.000    0.000    0.022    0.000 /usr/lib/python3.10/ast.py:33(parse)
        1    0.000    0.000    0.022    0.022 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py:463(run)
        7    0.000    0.000    0.022    0.003 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph.py:1562(python_code)
       27    0.000    0.000    0.021    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py:388(run_node)
        1    0.000    0.000    0.020    0.020 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/partitioners.py:285(_extract_fwd_bwd_modules)
        3    0.000    0.000    0.020    0.007 /home/xadupre/github/onnxscript/onnxscript/rewriter/__init__.py:28(rewrite)
       17    0.000    0.000    0.019    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/fx_onnx_interpreter.py:604(call_function)
      136    0.000    0.000    0.019    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1701(_output_from_cache_entry)
     12/9    0.001    0.000    0.019    0.002 {built-in method torch._C._nn.linear}
        7    0.000    0.000    0.019    0.003 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph.py:1639(_python_code)
      140    0.002    0.000    0.019    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1635(_get_output_tensor_from_cache_entry)
    26405    0.009    0.000    0.019    0.000 {method 'get' of 'dict' objects}
        7    0.002    0.000    0.019    0.003 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph.py:408(_gen_python_code)
     9915    0.010    0.000    0.019    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/registration.py:55(from_qualified_name)
     7575    0.006    0.000    0.018    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:123(_is_tensor_type)
8908/4120    0.009    0.000    0.018    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/node.py:883(map_aggregate)
        9    0.001    0.000    0.018    0.002 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:1498(wrap_tensor)
     12/9    0.000    0.000    0.018    0.002 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/nn/functional.py:1693(relu)
        1    0.000    0.000    0.018    0.018 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/guards.py:2107(__init__)
        9    0.001    0.000    0.018    0.002 {built-in method torch.relu}
       31    0.000    0.000    0.018    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:325(from_real_tensor)
      150    0.001    0.000    0.018    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1369(_cache_key)
        9    0.000    0.000    0.018    0.002 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/variables/torch.py:876(call_function)
       19    0.000    0.000    0.017    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/functional_utils.py:35(to_fun)
       19    0.000    0.000    0.017    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/functional_tensor.py:228(to_functional)
      796    0.000    0.000    0.017    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:168(is_attr_type)
       18    0.000    0.000    0.017    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_dynamo/utils.py:1700(wrap_fake_exception)
4085/3893    0.002    0.000    0.017    0.000 {built-in method builtins.next}
       31    0.001    0.000    0.017    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py:1588(__call__)
     4900    0.005    0.000    0.016    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/node.py:854(__setattr__)
  565/153    0.002    0.000    0.016    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1445(_prep_args_for_hash)
       82    0.000    0.000    0.016    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_guards.py:296(create)
      102    0.000    0.000    0.016    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/proxy.py:209(create_proxy)
  2434/12    0.001    0.000    0.016    0.001 /home/xadupre/github/onnxscript/onnxscript/ir/serde.py:94(wrapper)
    16577    0.006    0.000    0.014    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/onnx/_internal/fx/registration.py:45(from_name_parts)
        8    0.000    0.000    0.014    0.002 /home/xadupre/github/onnxscript/onnxscript/_legacy_ir/visitor.py:786(visit_model)
     2879    0.002    0.000    0.014    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/node.py:874(map_arg)
     1381    0.002    0.000    0.014    0.000 /home/xadupre/github/onnxscript/onnxscript/converter.py:451(_eval_constant_expr)
     8052    0.004    0.000    0.013    0.000 /home/xadupre/github/onnxscript/onnxscript/type_annotation.py:70(_remove_annotation)
        3    0.000    0.000    0.013    0.004 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph_module.py:908(print_readable)
        3    0.000    0.000    0.013    0.004 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/graph_module.py:297(_print_readable)
       18    0.000    0.000    0.013    0.001 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:2364(from_tensor)
      539    0.000    0.000    0.012    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_pytree.py:866(tree_flatten)
      238    0.002    0.000    0.012    0.000 /usr/lib/python3.10/inspect.py:932(findsource)
 1860/539    0.003    0.000    0.012    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/utils/_pytree.py:845(_tree_flatten_helper)
        1    0.000    0.000    0.012    0.012 /home/xadupre/github/onnxruntime/build/linux_cuda/Release/onnxruntime/capi/onnxruntime_inference_collection.py:406(__init__)
        1    0.012    0.012    0.012    0.012 /home/xadupre/github/onnxruntime/build/linux_cuda/Release/onnxruntime/capi/onnxruntime_inference_collection.py:484(_create_inference_session)
       61    0.000    0.000    0.012    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py:593(track_tensor_tree)
2155/2060    0.001    0.000    0.011    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/node.py:889(<listcomp>)
        8    0.000    0.000    0.011    0.001 /home/xadupre/github/onnxscript/onnxscript/_legacy_ir/visitor.py:646(visit_graph)
      495    0.002    0.000    0.011    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/fx/node.py:367(prepend)
       31    0.001    0.000    0.011    0.000 /home/xadupre/vv/this/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py:687(meta_tensor)
done.

Benchmark exported models with ORT

def benchmark(shape):
    data = []
    data_mem_first_run = []
    data_mem_run = []
    confs = list(
        itertools.product(
            export_functions,
            ["CPU", "CUDA"],
        )
    )
    loop = tqdm(confs)
    print(f"number of experiments: {len(loop)}")
    for export_fct, p in loop:
        name = export_fct.__name__.replace("get_torch_", "")
        obs = {}  # system_info()
        obs["name"] = name
        obs["compute"] = p
        obs["export"] = name

        model, input_tensor = create_model_and_input()
        if p == "CUDA":
            if not has_cuda:
                continue
            model = model.cuda()
            input_tensor = input_tensor.cuda()
        try:
            exported_model = export_fct(model, input_tensor)
        except torch._dynamo.exc.BackendCompilerFailed as e:
            # Triton only supports devices of CUDA Capability >= 7.0,
            # but your device is of CUDA capability 6.1
            obs["error"] = str(e)
            data.append(obs)
            continue

        def call_model(
            export_fct=export_fct,
            exported_model=exported_model,
            input_tensor=input_tensor,
        ):
            res = exported_model(input_tensor).sum()
            return res

        stat = start_spying_on(cuda=1 if has_cuda else 0)
        try:
            call_model()
        except Exception as e:
            loop.set_description(f"ERROR-run: {name} {e}")
            obs.update({"error": e, "step": "load"})
            data.append(obs)
            stat.stop()
            continue
        memobs = flatten(stat.stop())
        memobs.update(obs)
        data_mem_first_run.append(memobs)

        # memory consumption
        stat = start_spying_on(cuda=1 if has_cuda else 0)
        for _ in range(0, script_args.warmup):
            call_model()
        memobs = flatten(stat.stop())
        memobs.update(obs)
        data_mem_run.append(memobs)

        obs.update(
            measure_time(
                call_model,
                max_time=script_args.maxtime,
                repeat=script_args.repeat,
                number=1,
            )
        )

        profile_function(name, call_model, with_args=False, suffix=f"run_{p}")

        loop.set_description(f"{obs['average']} {name} {p}")
        data.append(obs)
        del model
        del exported_model
        gc.collect()
        time.sleep(1)

    df = pandas.DataFrame(data)
    df.to_csv("plot_torch_dort_ort_time.csv", index=False)
    df.to_excel("plot_torch_dort_ort_time.xlsx", index=False)
    dfmemr = pandas.DataFrame(data_mem_run)
    dfmemr.to_csv("plot_torch_dort_ort_run_mem.csv", index=False)
    dfmemr.to_excel("plot_torch_dort_ort_run_mem.xlsx", index=False)
    dfmemfr = pandas.DataFrame(data_mem_first_run)
    dfmemfr.to_csv("plot_torch_dort_ort_first_run_mem.csv", index=False)
    dfmemfr.to_excel("plot_torch_dort_ort_first_run_mem.xlsx", index=False)
    return df, dfmemfr, dfmemr


df, dfmemfr, dfmemr = benchmark(list(input_tensor.shape))
print(df)
  0%|          | 0/6 [00:00<?, ?it/s]number of experiments: 6

0.004093346852001829 eager CPU:   0%|          | 0/6 [00:00<?, ?it/s]
0.004093346852001829 eager CPU:  17%|█▋        | 1/6 [00:02<00:11,  2.21s/it]
0.0005163561082314054 eager CUDA:  17%|█▋        | 1/6 [00:02<00:11,  2.21s/it]
0.0005163561082314054 eager CUDA:  33%|███▎      | 2/6 [00:04<00:08,  2.08s/it]
0.0038502874074900456 default CPU:  33%|███▎      | 2/6 [00:04<00:08,  2.08s/it]
0.0038502874074900456 default CPU:  50%|█████     | 3/6 [00:06<00:06,  2.15s/it]
0.0005493982021480256 default CUDA:  50%|█████     | 3/6 [00:08<00:06,  2.15s/it]
0.0005493982021480256 default CUDA:  67%|██████▋   | 4/6 [00:09<00:05,  2.56s/it]
0.00044521257255908885 dort CPU:  67%|██████▋   | 4/6 [00:10<00:05,  2.56s/it]
0.00044521257255908885 dort CPU:  83%|████████▎ | 5/6 [00:12<00:02,  2.51s/it]
0.0007605620370735845 dort CUDA:  83%|████████▎ | 5/6 [00:12<00:02,  2.51s/it]
0.0007605620370735845 dort CUDA: 100%|██████████| 6/6 [00:14<00:00,  2.43s/it]
0.0007605620370735845 dort CUDA: 100%|██████████| 6/6 [00:14<00:00,  2.39s/it]
      name compute   export   average  deviation  min_exec  max_exec  repeat  number     ttime  context_size  warmup_time
0    eager     CPU    eager  0.004093   0.000377  0.002805  0.004365       1    27.0  0.110520            64     0.005876
1    eager    CUDA    eager  0.000516   0.000222  0.000327  0.000941       1   231.0  0.119278            64     0.001568
2  default     CPU  default  0.003850   0.000217  0.003427  0.004023       1    27.0  0.103958            64     0.002103
3  default    CUDA  default  0.000549   0.000112  0.000413  0.000896       1   183.0  0.100540            64     0.001423
4     dort     CPU     dort  0.000445   0.000114  0.000368  0.000974       1   255.0  0.113529            64     0.001761
5     dort    CUDA     dort  0.000761   0.000231  0.000635  0.001442       1   135.0  0.102676            64     0.002086

Other view

def view_time(df, title, suffix="time"):
    piv = pandas.pivot_table(df, index="export", columns=["compute"], values="average")
    print(piv)
    piv.to_csv(f"plot_torch_dort_{suffix}_compute.csv")
    piv.to_excel(f"plot_torch_dort_{suffix}_compute.xlsx")

    piv_cpu = pandas.pivot_table(
        df[df.compute == "CPU"],
        index="export",
        columns=["compute"],
        values="average",
    )

    fig, ax = plt.subplots(1, 2, figsize=(12, 4))
    fig.suptitle(title)
    piv_cpu.plot.barh(ax=ax[0], title="CPU", logx=True)

    if has_cuda:
        piv_gpu = pandas.pivot_table(
            df[df.compute == "CUDA"],
            index="export",
            columns=["compute"],
            values="average",
        )
        piv_gpu.plot.barh(ax=ax[1], title="CUDA", logx=True)

    fig.tight_layout()
    fig.savefig(f"plot_torch_dort_{suffix}.png")
    return ax


view_time(df, "Compares processing time on backends")
Compares processing time on backends, CPU, CUDA
compute       CPU      CUDA
export
default  0.003850  0.000549
dort     0.000445  0.000761
eager    0.004093  0.000516

array([<Axes: title={'center': 'CPU'}, ylabel='export'>,
       <Axes: title={'center': 'CUDA'}, ylabel='export'>], dtype=object)

Memory First Running Time (ORT)

for compute in ["CPU", "CUDA"]:
    if not has_cuda and compute == "CUDA":
        continue
    ax = memory_peak_plot(
        dfmemfr[dfmemfr.compute == compute],
        ("export",),
        suptitle=f"Memory Consumption of backend, first running time"
        f"\nrunning on {compute}",
        bars=[model_size * i / 2**20 for i in range(1, 3)],
        figsize=(18, 6),
    )
    get_figure(ax).savefig(f"plot_torch_dort_first_run_mem_{compute}.png")
  • Memory Consumption of backend, first running time running on CPU, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)
  • Memory Consumption of backend, first running time running on CUDA, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)

Memory Running Time (ORT)

for compute in ["CPU", "CUDA"]:
    if not has_cuda and compute == "CUDA":
        continue
    ax = memory_peak_plot(
        dfmemr[dfmemr.compute == compute],
        ("export",),
        suptitle=f"Memory Consumption of backens, running time\nrunning on {compute}",
        bars=[model_size * i / 2**20 for i in range(1, 3)],
        figsize=(18, 6),
    )
    get_figure(ax).savefig(f"plot_torch_dort_run_mem_{compute}.png")
  • Memory Consumption of backens, running time running on CPU, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)
  • Memory Consumption of backens, running time running on CUDA, Memory peak (Mb), Memory peak - memory begin (Mb), Memory average - memory begin (Mb), GPU Memory peak (Mb), GPU Memory peak - memory begin (Mb), GPU Memory average - memory begin (Mb)

Total running time of the script: (0 minutes 41.937 seconds)

Gallery generated by Sphinx-Gallery