EmergentOrder / onnx-scala

An ONNX (Open Neural Network eXchange) API and backend for typeful, functional deep learning and classical machine learning in Scala 3
GNU Affero General Public License v3.0
137 stars 8 forks source link

Accessing multiple outputs #367

Open neillbyrne opened 1 year ago

neillbyrne commented 1 year ago

Hi, firstly thanks for this application, its much appreciated

Im struggling to run an onnx model which uses more than a single output

image

I can get the first output values logits which has a dimension of (batch_size 86 1) fine with the following code, where I am using a single sample example with an input dimension of (batch_size * num_words), the last dimension being an array of word tokens ids (integers) which in this case is 3 for this example but is dynamic

val modelBytes = Files.readAllBytes(Paths.get(onnxPath))
val model = new ORTModelBackend(modelBytes)
val data =  Array.fill(1 * 3){10L}
val shape = 1 #: 3 #: SNil
val tsr = Tensor(data, shape)
val out = model.fullModel[Float,
        "logits",
        "batch_size" ##: "n_labels" ##: TSNil,
        1 #: 86 #: SNil](Tuple(tsr))

val output: Array[Float] = out.data.unsafeRunSync()
println(output.mkString(", "))
\\ prints 0.05673475, -3.656492, -1.9948144, -1.2033991, -0.457706, 0.15737924,

But the second output alphas doesn't seem accessible, it has dimensions (batch_size 86 num_words) or 1 86 3 for this example. I have tried running the model again to get the alphas, but it doesn't seem to work

val out = model.fullModel[Float,
    "alphas",
    "batch_size" ##: "n_labels" ##: "n_words" ##: TSNil,
    1 #: 86 #: 3 #: SNil](Tuple(tsr))

It raises a bunch of errors (which I suspect are due to incorrect shape in the spec)

java.lang.IllegalArgumentException: requirement failed

    at scala.Predef$.require(Predef.scala:324)
    at org.emergentorder.onnx.backends.ORTOperatorBackend.$anonfun$3(ORTOperatorBackend.scala:68)
    at flatten @ org.emergentorder.onnx.backends.ORTModelBackend.inputTensors$1$$anonfun$1(ORTModelBackend.scala:66)
    at blocking @ org.emergentorder.onnx.backends.ORTOperatorBackend.runModel(ORTOperatorBackend.scala:63)
    at make @ org.emergentorder.onnx.backends.ORTModelBackend.fullModel(ORTModelBackend.scala:75)
    at make @ org.emergentorder.onnx.backends.ORTModelBackend.fullModel(ORTModelBackend.scala:75)
    at use @ org.emergentorder.onnx.backends.ORTModelBackend.fullModel(ORTModelBackend.scala:83)
    at flatMap @ org.emergentorder.onnx.backends.ORTOperatorBackend.runModel(ORTOperatorBackend.scala:80)
    at flatten @ org.emergentorder.onnx.backends.ORTModelBackend.inputTensors$1$$anonfun$1(ORTModelBackend.scala:66)

I guess the most confusing part (besides this being my first foray into Scala!) is that the second parameter in the specification (which I assumed was the name of the output) seems to accept any string and access the first output logits, e.g.

val out = model.fullModel[Float,
        "what am I",
        "batch_size" ##: "n_labels" ##: TSNil,
        1 #: 86 #: SNil](Tuple(tsr))

works fine and returns the logits output

Many thanks in advance!

EmergentOrder commented 1 year ago

Hi, and thanks for the interest / feedback. As you correctly noted, this is currently hard-coded to return only the first output.

The type parameter you tried to use to select the output is just a tensor type denotation, which can be anything, but the type of which can be enforced at compile-time.

In the short term, I'll use the tensor type denotation to try to grab the output with the matching name, and if it doesn't exist fall back to the first output. So you will still need to run things once per output you want, but at least you can select an output other than the first.

Longer-term, it would make sense to add a way to return multiple outputs. I'll just need to consider exactly how best to do that while maintaining type safety.

neillbyrne commented 1 year ago

Sounds great @EmergentOrder . Many thanks for all your work!, feel free to close the issue