ONNX Support

[davideboschetto]

Hi there.

I see no discussions about ONNX, an open source framework to share models among different deep learning libraries. I don't know if this is better implemented at the Keras level or at the backend level (CNTK already supports it), but I think it's worth discussing!

[fchollet]

As of now, ONNX is a talking point used by large companies, that doesn't seem to address an existing problem and doesn't seem to have actual users. This may change in the future, but this is the current situation.

Keras takes a very programatic approach to development, implementing features that match user needs rather than abstract PR-driven strategies. If ONNX becomes significant in the future, we will add support for it. But while its purpose it to serve as a corporate PR talking point, we will not invest development efforts into it.

I encourage you to look at NNEF and OpenVX. Should Keras support them? Why? Why not? Is ONNX more relevant than these?

I would also point out that there is an existing NN exchange format that works across TF, CNTK, Theano, browsers, the JVM, and even MXNet to an extent: Keras savefiles. Importantly, Keras savefiles already have lots of users.

[hrisheekeshr]

Now that we have this, https://cloudblogs.microsoft.com/opensource/2020/01/21/microsoft-onnx-open-source-optimizations-transformer-inference-gpu-cpu/

Can we reopen this issue?

[Ark-kun]

there is an existing NN exchange format that works across TF, CNTK, Theano, browsers, the JVM, and even MXNet to an extent: Keras savefiles. Importantly, Keras savefiles already have lots of users.

Is there any documentation on how to use Keras savefiles in the latest version of Keras?

[davideboschetto]

I think reopening the issue might be the right choice!
ONNX provided this converter, in the meantime: https://github.com/onnx/tensorflow-onnx
It only works with saved_model or .meta, though (usual drawbacks of freezing the graph and exporting from there)

[innat]

cc @sachinprasadhs
This should be reopened.

[mattdangerw]

@innat can you say more about what you had in mind?

Something that goes directly keras.ops to the onnx format?

And how would that compare to going to onnx through a backend? E.g. set the torch bachend, use torch onnx tools to export to onnx?

[innat]

@mattdangerw

Something that goes directly keras.ops to the onnx format?

Sounds good. Here is jax2xla, a numpy backedn onnx interpreter, google/jaxonnxruntime.

And how would that compare to going to onnx through a backend? E.g. set the torch bachend, use torch onnx tools to export to onnx?

I think this suits more. However,

• If it is ensured that a jax-backend written keras mdoel can be run out of the box with torch-backend, then setting torch backend, torch.onnx can be used to export.
• AFAIK, there is no direct conversion from jax2onnx. To achieve this one can first convert jax2tf and then tf2onnx

[qlzh727]

Thanks all for the inputs.

I think the most approvable way to do this so far is either using TF savedmodel to ONNX, or torch backend (less mature on Keras side) to ONNX. For any Keras model, user probably can take following workflow:

1. build/train/test the Keras model in any backend.
2. Save the model in keras format.
3. Reload the keras model back with TF backend.
4. Export the keras model with TF backend to TF savedmodel.
5. Convert to ONNX.

Since Keras is a high level framework, we could also choose to directly export ONNX, but this will take signaficant effort to implement all the ops mapping, and I think leverage TF to ONNX is probably the most low cost approach at this moment.

Keras side probably should provide a guide/example for how to do this.

I believe this feature will be important since it opens up the way to user to leverage more downstream system like Nvidia TensorRT for inference, etc.

@fchollet for more inputs.

[qlzh727]

Discussed this in the team meeting.

In general we will leverage each backend for their own way to export to ONNX, eg TF has its own way to convert tf savedmodel to ONNX, same for JAX and pytorch. Keras will not support direct export to ONNX at this stage (eg implement ONNX ops set). We will provide documentation, and potential APIs to support user to convert to ONNX.

[qlzh727]

Adding Neel since this is related to saving/exporting.

[jojje]

Hi, any movement on this?
I'm considering keras again after the announcement that it once again supports pytorch after a very long hiatus of being a TF-only offering. But not being able to export my trained models to the now almost defacto inference formats for apps (ONNX) is currently a show-stopper for me.

[sachinprasadhs]

@jojje , Currently we have 2 items regarding ONNX in our roadmap mentioned here #19519

[guiyrt]

Hi, I would like to contribute towards this. From what I understood from this issue, the idea is to use each backend's existing onnx conversion method (at least for now).
The most straightforward is pytorch, we can just use torch.onnx. For tensorflow we need to introduce the tf2onnx dependency. For JAX, there is no official onnx conversion method, but the common workaround is first to convert to SavedModel using jax2tf from jax.experimental, and then using tf2onnx.
Is this what you had in mind?

[fchollet]

Is this what you had in mind?

Yes, exactly. Each backend can use its own mechanism for ONNX conversion. The public API would be via model.export, to which we would add an export_format argument. You are more than welcome to start a PR, and we'll guide you!

[guiyrt]

A few notes so far:

• I did some tests exporting keras.applications.vgg16 and the export worked fine for tensorflow and torch, but I couldn't make it work for jax
• There are a lot of differences arguments-wise with tf2onnx and torch.onnx, so it's difficult to have an export function that is backend agnostic (the main difference I found was in model input signature). For torch, you provide a dummy tensor with the expected shape, and with tensorflow you pass a tf.TensorSpec object describing the input. See tf2onnx.convert.from_keras and torch.onnx.export.

I created tests comparing the output of the original and exported models using onnxruntime.InferenceSession, and passed for both tf and torch. My major doubt is regarding the public API. I implemented it in model.export, as suggested, and updated the format argument to accept onnx (so far only tf_saved_model was accepted). This means we need to get onnx-related args here, and I'm not sure which is the best way for the user to do so. They can't be explicit unless we do the mapping ourselves (not trivial), so having something like **kwargs or a onnx_args dict decided by the user for the chosen backend might be the easier route. But that means some backend-specific logic and tests.

Let me know what you think, but it might be easier if I publish my dev branch and we tackle it from there.

[james77777778]

After a long run, we can use Model.export(filepath, format="onnx") to export ONNX models from TF, JAX, and Torch backends. Please let me know if you encounter any issues.