[WIP] ONNXExport.jl - Export Julia functions to ONNX

Package Announcements
, ,
1

Hello,

What

ONNXExport.jl is a package for exporting Julia functions as ONNX models. It is still an unregistered WIP but has a decent coverage of Julia, NNlib, MLUtils, and Lux functions. See the README for detailed coverage and all limitations.

Why

There are existing packages such as ONNX.jl and ONNXNaiveNASflux.jl that produce ONNX files. However, the former seems inactive and the latter is made for exporting Flux models. I needed to export Lux models. Another option would be to serialize a Lux model to the TensorFlow SavedModel format and then convert to ONNX, but that is limited to static dimensions. I needed symbolic dimensions/dynamic shapes.

Usage

The package exposes an ONNXExport.save function which traces a provided function and saves the resulting ONNX model to file.

using ONNXExport

f(x, y) = x .+ y .- 3
ONNXExport.save("model.onnx", f, rand(Float32, 3, 4), rand(Float32, 3))

We can export Lux models by wrapping Lux.apply in a function f(x). Symbolic dimensions, e.g., unknown batch size, is supported but largely untested.

using ONNXExport, Lux, Random

model = Chain(Dense(16 => 8, relu), Dense(8 => 2))

rng = Random.default_rng()
ps, st = Lux.setup(rng, model)
st = Lux.testmode(st)

f(x) = first(Lux.apply(model, x, ps, st))
ONNXExport.save("model.onnx", f, TypeInfo(16, :N))

The last arguments to ONNXExport.save can be example inputs from which types and sizes are inferred. Alternatively, this information can be provided explicitly using TypeInfo.

What’s Next

Before I spend more time on this, I would like to know:

  • Is this something people find useful?
  • Will this be replaced by something better soon?

I suspect Reactant will eventually support dynamic shapes/symbolic dimensions, and then serialization directly to ONNX could be possible. (Perhaps I should have spent my time doing that as a contribution to Reactant rather than writing yet another tracer.)

In any case, I created this to solve problems I had. Let me know if this could solve problems you have as well. The package is unregistered and the documentation takes the form of the README and docstrings. The external API is simple enough to not change drastically, but the internals are very liquid at the moment.

GitHub: ONNXExport.jl

Martin Larsson

2 Likes