Intermediate results with onnxruntime

Example Intermediate results with (ONNX) ReferenceEvaluator demonstrated how to run a python runtime on a model but it may very slow sometimes and it could show some discrepancies if the only provider is not CPU. Let’s use OnnxruntimeEvaluator. It splits the model into node and runs them independently until it succeeds or fails. This class converts every node into model based on the types discovered during the execution. It relies on InferenceSessionForTorch or InferenceSessionForNumpy for the execution. This example uses torch tensor and bfloat16.

A failing model

The issue here is a an operator Cast trying to convert a result into a non-existing type.

import onnx
import onnx.helper as oh
import torch
import onnxruntime
from onnx_diagnostic import doc
from onnx_diagnostic.ext_test_case import has_cuda
from onnx_diagnostic.helpers.onnx_helper import from_array_extended
from onnx_diagnostic.reference import OnnxruntimeEvaluator

TBFLOAT16 = onnx.TensorProto.BFLOAT16

model = oh.make_model(
    oh.make_graph(
        [
            oh.make_node("Mul", ["X", "Y"], ["xy"], name="n0"),
            oh.make_node("Sigmoid", ["xy"], ["sy"], name="n1"),
            oh.make_node("Add", ["sy", "one"], ["C"], name="n2"),
            oh.make_node("Cast", ["C"], ["X999"], to=999, name="failing"),
            oh.make_node("CastLike", ["X999", "Y"], ["Z"], name="n4"),
        ],
        "-nd-",
        [
            oh.make_tensor_value_info("X", TBFLOAT16, ["a", "b", "c"]),
            oh.make_tensor_value_info("Y", TBFLOAT16, ["a", "b", "c"]),
        ],
        [oh.make_tensor_value_info("Z", TBFLOAT16, ["a", "b", "c"])],
        [from_array_extended(torch.tensor([1], dtype=torch.bfloat16), name="one")],
    ),
    opset_imports=[oh.make_opsetid("", 18)],
    ir_version=9,
)

We check it is failing.

try:
    onnxruntime.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
except onnxruntime.capi.onnxruntime_pybind11_state.Fail as e:
    print(e)

[ONNXRuntimeError] : 1 : FAIL : Node (failing) Op (Cast) [TypeInferenceError] Attribute to does not specify a valid type in .

OnnxruntimeEvaluator

This class extends onnx.reference.ReferenceEvaluator with operators outside the standard but defined by onnxruntime. verbose=10 tells the class to print as much as possible, verbose=0 prints nothing. Intermediate values for more or less verbosity.

ref = OnnxruntimeEvaluator(model, verbose=10)
feeds = dict(
    X=torch.rand((3, 4), dtype=torch.bfloat16), Y=torch.rand((3, 4), dtype=torch.bfloat16)
)
try:
    ref.run(None, feeds)
except Exception as e:
    print("ERROR", type(e), e)

 +C one: bfloat16:(1,):[1.0]
 +I X: D-1:torch.bfloat16:torch.Size([3, 4]):0.34765625,0.453125,0.23828125,0.53515625,0.390625,0.57421875,0.47265625,0.4140625,0.6640625,0.7109375...
 +I Y: D-1:torch.bfloat16:torch.Size([3, 4]):0.2890625,0.96875,0.51171875,0.03125,0.85546875,0.3046875,0.20703125,0.99609375,0.1875,0.2578125...
Mul(X, Y) -> xy
ERROR <class 'onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented'> [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for Mul(14) node with name 'n0'

onnxruntime may not support bfloat16 on CPU. See onnxruntime kernels.

if has_cuda():
    ref = OnnxruntimeEvaluator(model, providers="cuda", verbose=10)
    feeds = dict(
        X=torch.rand((3, 4), dtype=torch.bfloat16), Y=torch.rand((3, 4), dtype=torch.bfloat16)
    )
    try:
        ref.run(None, feeds)
    except Exception as e:
        print("ERROR", type(e), e)

 +C one: bfloat16:(1,):[1.0]
 +I X: D-1:torch.bfloat16:torch.Size([3, 4]):0.31640625,0.64453125,0.89453125,0.4140625,0.7265625,0.30859375,0.390625,0.6640625,0.5859375,0.140625...
 +I Y: D-1:torch.bfloat16:torch.Size([3, 4]):0.24609375,0.01953125,0.37890625,0.59765625,0.3671875,0.2890625,0.5625,0.37890625,0.03515625,0.05078125...
Mul(X, Y) -> xy
 + xy: D-1:torch.bfloat16:torch.Size([3, 4]):0.07763671875,0.0125732421875,0.33984375,0.2470703125,0.267578125,0.08935546875,0.2197265625,0.251953125,0.0206298828125,0.00714111328125...
Sigmoid(xy) -> sy
 + sy: D-1:torch.bfloat16:torch.Size([3, 4]):0.51953125,0.50390625,0.5859375,0.5625,0.56640625,0.5234375,0.5546875,0.5625,0.50390625,0.50390625...
Add(sy, one) -> C
 + C: bfloat16:(3, 4):1.515625,1.5,1.5859375,1.5625,1.5625,1.5234375,1.5546875,1.5625,1.5,1.5...
Cast(C) -> X999
ERROR <class 'RuntimeError'> Unable to infer a session with inputs
#1[A16r2]
due to [ONNXRuntimeError] : 1 : FAIL : Node (failing) Op (Cast) [TypeInferenceError] Attribute to does not specify a valid type in .
opset: domain='' version=18
input: name='C' type=dtype('float32') shape=[3, 4]
Cast(C, to=999) -> X999
output: name='X999' type='NOTENSOR' shape=None

We can see it run until it reaches Cast and stops. The error message is not always obvious to interpret. It gets improved every time from time to time. This runtime is useful when it fails for a numerical reason. It is possible to insert prints in the python code to print more information or debug if needed.

doc.plot_legend("onnxruntime\nrunning\nstep by step", "OnnxruntimeEvaluator", "lightgrey")

Related examples

Intermediate results with (ONNX) ReferenceEvaluator

Intermediate results with (ONNX) ReferenceEvaluator

Find where a model is failing by running submodels

Find where a model is failing by running submodels

Find and fix an export issue due to dynamic shapes

Find and fix an export issue due to dynamic shapes