1968939 - Add an API to optimize ONNX models

Status: NEW

(Core :: Machine Learning: On Device, enhancement, P3)

Description

[Paul Adenot (:padenot)]

It's exposed via MLUtils in js, and allows optimizing models after having them downloaded. Then when inferencing, the optimized model is used, saving a large amount of time on the critical path, reducing end-to-end latency for the user.

Comment 1

[Tarek Ziadé (:tarek)]

This change will come in several components:

new C++ APIs surfaced in JS:

• splitOnnxFile : a function in MLUtils that gets an ONNX model and splits it into a separate graph and data, bonus point if we can make this one streamable so we don't load the whole model in memory.
• compileGraph : a function in MLUtils that gets an ONNX graph and returns a compiled graph

From there, ModelHub can decide in the main process at download time if it wants to compile on the fly the model and split data.

If it does, it will store for each model.onnx file:

• model.onnx : the compiled graph, stripped of the weight data
• model.onnx_data : the weights

And we'll use those files when running inference (the WASM runtime already supports that)

Comment 2

[Tarek Ziadé (:tarek)]

Notice that we will also need to extend the onnx-native backend call, by feeding the data in session_options.externalData

We should also deactivate in the runtime the optimization step and make the assumption it's done before

Comment 3

[Tarek Ziadé (:tarek)]

For reference, our Python script that splits graph and weights: https://searchfox.org/mozilla-central/source/toolkit/components/ml/tools/convert_to_external_data.py

Comment 4

[Paul Adenot (:padenot)]

Attachment: Bug 1968939 - Updater onnx api header to 1.22. r?tarek

Comment 5

[Paul Adenot (:padenot)]

Attachment: Bug 1968939 - Move utility functions in InferenceSession to a dedicated file. r?tarek

Comment 6

[Paul Adenot (:padenot)]

Attachment: Bug 1968939 - Expose a function to js to compile an ONNX model with optimization. r?tarek

Comment 7

[Paul Adenot (:padenot)]

Clarification: this is for now blocked on https://github.com/huggingface/transformers.js/pull/1382, that we need because it pulls in a new onnxruntime update, that we need to not risk ABI breakage.

Comment 8

[Greg Tatum [:gregtatum]]

I'm simplifying the dependency tree a bit, as I'm finding it confusing what works needs doing here.