How to include an ONNX model into onnxscript script?

[kubaraczkowski]

Hi onnnxscript people,

Onnxscript is great, I love it, thanks!
I currently use it e.g. for preprocessing tasks, but would want to include the 'real' model in the flow as well.
More, I would e.g. like to include 3 models as an assembly in an onnxscript.

Until now I didn't find a way (tutorial, documentation) on how to do it. It's perhaps slightly related to #952 which converts the ONNX model to a script, I guess to also be able to 'merge' it with other scripts (and tweak the whole thing).

Therefore, in essence - is it possible to include a 'compiled' (proto) ONNX model as a part of an onnxscript?

Thanks!

[justinchuby]

Cc @gramalingam

[gramalingam]

Agree, it would be great to have this feature. Unfortunately, onnxscript doesn't have it yet.

Ideally, we should have a wrapper function that allows us to treat an existing ONNX model as a function, and be able to call it from an onnxscript function. Would be great to add this feature.

[kubaraczkowski]

Yes!
My workaround is the following:

• write the preprocessing and postprocessing functions
• merge them with the models using onnx.compose.merge_models. Not as trivial since names often clash, so requires adding prefixes everywhere
• run using onnxruntime

Works, but misses the fantastic debuggability of onnxscript!

Looking forward to this happening :)

[titaiwangms]

Feel free to reopen if the question is not answered.

[kubaraczkowski]

The workaround mentioned earlier works OK, but is REALLY cumbersome, against the ease-of-use of onnxscript.

With anything but a trivial model graph it becomes a pain to manage the manual merging of all the models, manually linking the inputs and outputs.
Additionally you need to prefix all the onnxscript-generated models as they get generated from the same "empty" starting point.

So although this is possible, a function wrapper that @gramalingam suggested would be oh-so-much-better. Any chance it could be put on the roadmap?

[justinchuby]

Thank you for your feedback! I will take a closer look at this.

[kubaraczkowski]

Any hints on how to implement this?

The workaround with using merging works, but gets really cumbersome with more than one "file-based" model.
Additionally the most awesome feature (ok, one of, there are many in onnxscript!) of interactive model debugging is unavailable this way.

Thanks once again for the project!

[NoRaincheck]

I'm interested in this as well. My gut feel is that maybe if OnnxFunction was more public facing?

my_model = onnx_ir.load("my_model.onnx")
my_model_as_onnx_function = OnnxFuction(..., irfun=<something_here>(my_model), ...)

It obviously breaks a lot of assumptions in OnnxFunction though.

---

Edit: I suppose something like this works, though its not documented. May also be fragile depending on how the ModelProto to FunctionProto conversion goes

import numpy as np
import onnx
import onnx_ir
import onnxruntime
from onnxscript import FLOAT, script
from onnxscript import opset15 as op
from onnxscript.values import Opset

# A dummy opset used for model-local functions
local = Opset("local", 1)


@script(local, default_opset=op)
def diff_square(x, y):
    diff = x - y
    return diff * diff


@script(local)
def sum(z):
    return op.ReduceSum(z, keepdims=1)


@script()
def l2norm(x: FLOAT["N"], y: FLOAT["N"]) -> FLOAT[1]:  # noqa: F821
    return op.Sqrt(sum(diff_square(x, y)))


@script()
def l2norm2(x: FLOAT["N"], y: FLOAT["N"]) -> FLOAT[1]:  # noqa: F821
    return op.Sqrt(local.sum(local.diff_square(x, y)))


def convert_model_to_function(model: onnx.ModelProto, domain, name) -> onnx.FunctionProto:
    model_ir = onnx_ir.serde.deserialize_model(model)
    function_ir = onnx_ir.Function(domain=domain, name=name, graph=model_ir.graph, attributes={})
    return onnx_ir.to_proto(function_ir)


a_non_onnxscript_function = convert_model_to_function(sum.to_model_proto(), "local", "sum")


model = l2norm2.to_model_proto(functions=[a_non_onnxscript_function, diff_square])
print(onnx.printer.to_text(model))


session = onnxruntime.InferenceSession(model.SerializeToString())
print(
    session.run(
        None, {"x": np.array([1.0, 2.0, 3.0]).astype(np.float32), "y": np.array([4.0, 5.0, 6.0]).astype(np.float32)}
    )
)