ONNX import/export

[DrChainsaw]

I’m willing to put some effort in the ONNX story if there exist some interest.

TBH I don’t really know what fastAI is about so if there is some special meaning to ONNX + fastAI then please let me know.

Non exhaustive rundown of the current status Flux is that @opus111 has created BaseOnnx and I have made a branch of ONNXmutable which makes use of it (replacing ONNX.jl as the source of protos).

Afaik, ONNXmutable is fully functional and verified and the main gripe I have with it is that it is a bit of a monolith. It depends on NaiveNASflux for the model DAG which in turn has the somewhat big dependencies Flux (arguably not a big deal), LightGraphs and MetaGraphs (both which could be removed though with a little effort), JuMP and Cbc. ONNXmutable also has onnx and onnxruntime as test-only dependencies as well as PyCall and Conda to be able to use and depend on them.

As such, I’m mainly focused on breaking down that monolith into smaller and more reusable components which don’t force all those dependencies on to people. Please let me know if this is not what you think is needed here.

Here are some of my thoughts:

Import

In addition to the primitives discussed below, one needs some kind of runnable representation of the computation graph.

Julias fantastic autodiff capabilities have made it obsolete to have a special DAG format for the models (e.g. Tensorflow) as you can just write the DAG as a normal Julia function. While it is indeed possible to translate an ONNX model into a Julia expression which evaluates to a function representing the model (this is the approach taken by ONNX.jl) it has the limitation that the expression (at least in practice) has to be evaluated at the top level.

Flux today does not have a typical DAG. Flux’s built in Chain can represent many DAGs through the usage of SkipConnection but I don’t think it can represent any DAG (without user written functions).

I think that the method used in ONNXmutable can be generalized without too much effort to work for any typical DAG and perhaps this is something which is useful to put in BasicOnnx. I opened an issue about it here.

If there is interest, I think I can make a functor compatible import-onnxmodel-as-a-function macro in BaseOnnx. Drawback is perhaps that it might give people a janky experience due to it only working from the top level, but perhaps this can be solved by generating a better error message than “incorrect world age”.

To reach the end goal of having a useable package, there still needs to be a DAG format to use. One option is that I extract the non-mutation stuff from NaiveNASlib into some AbstractMlDag package. It is already separated from the mutation stuff so it is trivial to move to a separate package.

Export

The method in ONNXmutable uses dispatch for tracing and I think this is good enough for most typical/traditional DL models. It will however fail if it encounters

1. (Non primitive) functions with type annotations (ie function model(x::AbstractArray) )
2. Non-function expressions, e.g. if/else and for
   I think IRTools can be used without too much effort to circumvent 1, but I dread to think about how to make use of it for 2. Mjolnir seems to have the perfect abstraction but I don’t know if it is effectively maintained and ready for use.

For exporting, I’m not as certain that there exist some simple and universal enough solution which is worthwhile to put in BaseOnnx. Should we try to make a generic package or two for this or just mash it together with the framework specific stuff?

Primitives

Primitives are the functions which have a one to one correspondence with an operator defined in ONNX, for example Add, sin, Conv, RNN etc. In other words, this is the part which knows how to transform e.g. a Flux.Conv into an ONNX NodeProto and vice versa.

I don’t think this can be done without manually typing out the mapping for each OP so in any moderately well designed ONNX package this will be the by far biggest effort to create and maintain, especially considering opset versions. To me this makes it quite important that adding OPs can be done easily so that users of the package can contribute with the OPs they need or else it will be a thankless and soul crushing effort to support the whole spec.

To me, this makes it pretty useful to have a package with only the primitives to attract contributions. Obviously one package per ML framework is needed as the primitive package has to depend on the ML framework package (e.g. Flux, KNet, Lilith etc).

Here we can start with the modest set of primitives from ONNXmutable as a start of the Flux package.

Furthermore, import and export have basically nothing to do with each other and I don’t see a way to make use of import primitives for export and vice versa. This opens up for having separate packages for import and export primitives, but it is also kinda nice to be able to test the import functionality with the export functionality and vice versa. Thoughts on this?

Another thought is whether it makes sense to have a package with OPs from Base (e.g. Add, tan, sin, max, reshape etc)?

Testing

Testing numerical computations is always annoying. In ONNXmutable I used 1) test vectors from onnx and 2) comparison with output from onnxruntime. Nothing says onnxruntime is the golden standard ofc, but it seems like a lot of work is being put into it so I think it makes for a pretty good reference.

There are a couple of lines of code to set all of this up and perhaps this can also be made a package which others can make use of (e.g. primitive packages for other ML frameworks).

Package homes

The packages which are specific to a particular ML framework are best served to sit in the same org as their parent frameworks, right? What about the reminders?

Assuming the above, the reminders are basically BaseOnnx, the testing package, the Base OPs and maybe a generic export package or two if it makes sense to create it.

I guess that creating a new org in github is no effort, but perhaps it is good to try to put it into some more well known org like JuliaML.

Another option is to just give up on components which are not tied to any framework and perhaps just create another monolith. I think the advantage of this over what ONNXmutable offers today then is basically that one does not need to depend on JuMP and Cbc. I can't imagine that this is the reason why people state that julia has no functional ONNX import/export though.

[DrChainsaw]

Here is my takeaway from the meeting:

Splitting up ONNXmutable has lower priority than making sure it can load state of the art models. Is this correct? Quickest way to proceed here is to just test loading models and file issues (and preferably PRs 😊 ) when they fail.

If so, I can probably take the functionality from BaseOnnx and add it to ONNXmutable instead along with (protobuf 0.9 compliant) protos in order to get rid of the warning-generating dependency to ONNX.jl and release a 0.1 version very soon. If there is a need to split it up for reuse purposes this can always be done later of course.

I will add examples of loading a model and replacing the head as well as a pruning example here. Please let me know if there is anything else you’d like to see.

Load model and replace classification head (click me to expand)

Here is how to load a model and a few of the ways NaiveNASflux lets you explore the model

(ImportOnnxExample) pkg> add NaiveNASflux, MLDatasets, https://github.com/DrChainsaw/ONNXmutable.jl

julia> using ONNXmutable, NaiveNASflux, MLDatasets
# Lots of warnings from ONNX.jl :(

julia> resnetfile = download("https://github.com/onnx/models/raw/master/vision/classification/resnet/model/resnet18-v1-7.onnx");

julia> model = CompGraph(resnetfile);

julia> name.(vertices(model)) # Query names like this

julia> nout.(vertices(model)) # List output sizes like this 

Ok, lets scrape off the classification head and replace it with a new dense layer.

julia> vs = vertices(model);

julia> remove!(vs[end], RemoveStrategy(NoSizeChange())); # Normally we want NaiveNASflux to align the sizes between the input to the removed layer and the output, but in this case we explicitly want to change the output size without touching anything else

julia> insert!(vs[end-1], v -> mutable("newhead", Dense(nout(v), 10), v));

julia> newmodel = CompGraph(vs[begin], outputs(vs[end-1]));

julia> vertices(newmodel) |> last |> nout
10

julia> vertices(newmodel) |> last |> name
"newhead"

Unfortunately removing the very last layer is slightly more inconvenient than any other because 1) the CompGraph has a reference the output layer and 2) because NaiveNASflux will by default think that it needs to change the size of the previous layer to have the same output size (which is a pretty reasonable default assumption to be fair).

In normal cases one can just do remove! and insert! and it will change the model without having to create a new model.

Anyways, now one can just (re-)train the model as normal.

julia> Flux.adapt(T, x::Flux.Zeros) = x # Workaround for Flux issue # 1332

julia> Flux.adapt(T, x::Base.ReinterpretArray) = T(x)  # ONNX.jl turns TensorProtos into `ReinterpretedArray`s and Zygote does not like that. BaseOnnx makes them normal Array, so this is a temporary nuisance

julia> newmodel = newmodel |> cpu;  # Same as above (this should ofc be gpu if training on gpu)

julia> loss(x,y) = NaiveNASflux.Flux.Losses.logitcrossentropy(newmodel(x), y); # forgot to add Flux, but it is available through NaiveNASflux

julia> x,y = MLDatasets.CIFAR10.traindata();

julia> xpadded = zeros(Float32, 224,224, 3, 8); # Need to pad excessively to match size for this constructed use case

julia> xpadded[97:128,97:128,:,:] .= x[:,:,:,1:8];

julia> yhot = NaiveNASflux.Flux.onehotbatch(y[1:8], 0:9);

julia> pshead = NaiveNASflux.Flux.params(vertices(newmodel) |> last |> layer); # Only train the parameters of the last layer

julia> length(pshead)
2 

julia> NaiveNASflux.Flux.train!(loss, pshead, [(xpadded, yhot)], ADAM()) # I don’t have a GPU on this computer, so I’ll only do one batch

[DrChainsaw]

Here is the pruning use case

julia> model = CompGraph(resnetfile);

julia> numneurons(m) = mapreduce(nout, +, vertices(m)); # Not needed, just to show that something happened

julia> numneurons(model)
15531

julia> function pruning_metric(v, offs)
      val = neuron_value(v) # neuron_value defaults to magnitude of parameters along activation dimension
      ismissing(val) && return fill(offs, nout_org(v)) # Layers with no parameters return missing by default
      return val .- min(offs, 0.8*maximum(val)) # min is crude safeguard to prevent that models get size 0
      end

julia> allvals = mapreduce(neuron_value, vcat, vertices(model)) |> skipmissing |> collect;

julia> cutoff = partialsort(allvals, round(Int, 0.3*length(allvals))) # Cutoff is bigger than (approx) 30% of all values

julia> for v in vertices(model) # It is currently a limitation that one first must reduce the size of each individual vertex. If you want to use NaiveNASflux for pruning I think I can fix this limiation
      metric = pruning_metric(v, cutoff)
      nprune = sum(<(0), metric) - length(metric) + nout(v)
      nprune <= 0 && continue
      Δnout(v, -nprune)
      end

julia> Δoutputs(model, v -> pruning_metric(v, cutoff)); # Given that we have new sizes for all layers, which neurons do we decide to keep

julia> apply_mutation(model);

julia> model(ones(Float32, 224,224,3,2)) |> size # Model is still internally consistent
(1000, 2)

julia> numneurons(model) # 30% fewer neurons
10042

Now, this turned out to prune alot more than 30% of all parameters, so chances are the model accuracy suffered alot. It also seems like the layers with more parameters also tend to have lower magnitude of their parameters, so the strategy to have a global cutoff is probably not the best.

It should however be straight forward to change the above example to do a comparison per layer instead (i.e prune 30% of each layer instead of 30% in total) so I'll leave that exercise to the reader :).

Note that to train this model one needs to do that annoying (and soon to be fixed) dance with overloading a few Flux methods and mapping to cpu.

Obviously all the above examples could (and perhaps should) be wrapped in much nicer APIs. NaiveNASflux tries to be more of a library than a user-facing package and I have put flexibility before ease of use.

I'm thinking this can be used as a start to figure out what a nice looking API for the above use cases could look like.

[jeremiedb]

Maybe a naive question, but would it be possible to convert the model imported by ONNXmutable back into a regular Flux chain?

For example, following:

julia> model = CompGraph(resnetfile);

If the model could be recompose into a the stack of regular Flux building blocks (Conv, Dense, SkipConnexion), then I think it would make the manipulation of the model, such as for image transfer learning very easy and intuitive, just like what is done in the model-zoo tutorial: https://github.com/FluxML/model-zoo/blob/master/tutorials/transfer_learning/transfer_learning.jl.
The issue with the latter being that only VGG19 seems functional in Metalhead.

[DrChainsaw]

It certainly could, at least for the type of models which can be expressed as a chain. It is a bit of a mental exercise since the graph format in ONNX is very different from how (non-linear) graphs are expressed with chains. The current deserialization is just a very simple recursion through the graph and I don't think one can do the same to create a chain. Splitting ONNXmutable into multiple packages would however allow for creating a simpler importer which makes use of the same primitives, but it seems like the interest in this is somewhat low.

Another somewhat annoying but certainly overcomeable issue is that one probably needs to resort to wrapping layers in closures and that requires that some side mechanism to provide the parameters is provided.

I also think that the example looks easy because the graph is very linear. The same mechanism would not be so nice in the non-linear case. I guess that the "scrape off the classification head" usecase would be fine for almost all models though.

[DrChainsaw]

Btw, if the graph is linear then one can just do this:

julia> chain = Chain(layer.(vertices(compgraph)[2:end])...)

[jeremiedb]

Thanks a lot for the clarifications!
My general impression was that the ability to easily access typical pre-trained building blocks such as the VGG/ResNet models was a quite common need, so I've been a little surprised that it didn't appear as a high priority to the Flux ecosystem.
And on the NLP side

The Chain(layer.(vertices(compgraph)[2:end])...) works smoothly withh VGG16/19 models.
For ResNet, it seems to struggle:

graph = CompGraph("data/resnet34-v1-7.onnx")
julia> m = Chain(layer.(vertices(graph)[2:end])...)
ERROR: MethodError: no method matching layer(::NaiveNASlib.var"#225#226"{typeof(+)})
Closest candidates are:
  layer(::ONNXmutable.Flatten) at C:\Users\jerem\.julia\packages\ONNXmutable\kxK8z\src\deserialize\constraints.jl:179
  layer(::CompVertex) at C:\Users\jerem\.julia\packages\NaiveNASflux\0lGnm\src\vertex.jl:114
  layer(::NaiveNASflux.InputShapeVertex) at C:\Users\jerem\.julia\packages\NaiveNASflux\0lGnm\src\vertex.jl:18

Also, with the ResNets v2, it fails at the import step:

julia> graph = CompGraph("data/resnet34-v2-7.onnx")
ERROR: MethodError: no method matching (::ONNXmutable.var"#144#145")(::Dict{Symbol,Any})
Closest candidates are:
  #144(::Any, ::Any) at C:\Users\jerem\.julia\packages\ONNXmutable\kxK8z\src\deserialize\ops.jl:214
Stacktrace:
 [1] wrapfrom(::ONNXmutable.OnnxNode, ::ONNXmutable.OnnxNode, ::ONNXmutable.CompGraphBuilder, ::Symbol, ::Dict{Symbol,Any}) at C:\Users\jerem\.julia\packages\ONNXmutable\kxK8z\src\deserialize\combine.jl:47
...

Are the residual connexions the cause the difficulties here?
The ONNX models were taken here: https://github.com/onnx/models/tree/master/vision/classification/resnet/model

(let me know it'd be preferable that I open an issue the ONNXmutable repo or elsewhere)

[DrChainsaw]

Thanks for showing interest @jeremiedb, sorry for wall-of-texting you as a response :)

My general impression was that the ability to easily access typical pre-trained building blocks such as the VGG/ResNet models was a quite common need, so I've been a little surprised that it didn't appear as a high priority to the Flux ecosystem.

I think this has been stated a few times, but outside of the group maintaining this tracker it does not seem like there is an overwhelming need, judging by the number of posts on discourse and the influx of issues in ONNX.jl and ONNXmutable.jl. There could certainly be a chicken and egg problem here as people might silently turn back to python when they can't find a canonical ONNX package for any of the Julia ML frameworks.

I would love to discuss my proposal above some more. I'm however a bit hesistant to just go ahead and litter the general registry with ONNX packages without any wider support. Chances are that people will just go "hmmm, unkown single author and I don't understand the point of all of this so I'll just roll my own" if I proceed. I'm happy to do some work if we can work out a structure which we believe in. I will also be happy if the conclusion is that we should wait for some ONNX expert to catch interest instead :)

For ResNet, it seems to struggle:

Yes, sorry for using confusing home-made graph terminology here. A ResNet would not fall into the category of "linear DAGs" as the elementwise summations are nodes in the graph which take input from more than one other node.

The CompGraph basically uses the exact same graph representation as ONNX: Each node is an operation and its input edges describe from which other nodes in the graph it takes input and the output edges describe which other nodes consume the output. These are accessible through inputs(node) and outputs(node) respectively (I ended up using the name vertex instead of node in NaiveNASlib but afaik node and vertex are direct synonyms). With a CompGraph g in hand, one can recurse through the node using these methods and either g.inputs or g.outputs.

I can't say I have though alot about it, but creating an algorithm to transform an arbitrary graph using this representation to a chain does not seem like a trivial task. I'm sure the resnet can be done if one assumes that the only thing one needs to handle are things which fit the SkipConnection (e.g. by memoizing the inputs when there are more than one and wrap in a SkipConnection when you hit one of the memoized inputs again), but how would one do e.g. a ResNext?

If you have a suggestion for the above I'd be happy to accept a PR (or even implement it if you tell me how) in ONNXmutable to have something like Chain(g::GraphProto) or just load(modefile; as=Chain) while pending input for the structure proposal above.

Also, with the ResNets v2, it fails at the import step:

This is definitely an issue with ONNXmutable! Please file an issue and I'll look into it.

FYI the code which fails are the heuristics to combine multiple ONNX nodes into a single CompGraph node. One example of when one might want this are activation functions as ONNX always has them as separate nodes in the graph while Flux allows them to be inside the layers. This is by no means a necessary step, but I though it was nice to allow for things like CUDA optimizations as well as just generally trying to make sure that import -> export returns something which is as close to the same model as possible.

[darsnack]

There are a couple of related issues here too.

First, on the point of basic models/blocks like VGG/ResNet, we have a PR to Metalhead.jl that implements some of these natively in Flux. Ideally, I think we should skip ONNX completely and provide pre-trained models by just training those native Flux models (that was my intent anyways).

Second, there would still be a need beyond basic models for the ability to read in arbitrary ONNX models into Flux models. I agree that this would be very useful functionality, but like @DrChainsaw mentioned, it is a tricky problem. Any solution would probably need to be well-documented so someone can reference what to expect from the translation. That being said, I think this PR will be helpful for expressing more complex graphs. I've already indicated in one of the ML committers calls that the PR should get some attention and be merged.

[DrChainsaw]

Second, there would still be a need beyond basic models for the ability to read in arbitrary ONNX models into Flux models.

I'm not sure I would just blanket equate "Flux models" with "Chain". In my view, a CompGraph is also a Flux model, as is any julia function which makes use of building blocks from Flux (i.e what ONNX.jl tried to do).

That being said, I think this PR will be helpful for expressing more complex graphs.

I still don't fully understand the rationale behind that PR. Isn't the nice thing about the Chain that if your model is super simple so that it only consists of unary operations stringed (well chained actually) you can get rid of alot of complexity as you only need to store the layers as a tuple/array since the structure is implicit from the struct itself?

Once you deviate from this simple assumption, wouldn't it be better to just use the same format as what is more or less tried and true when it comes to DAGs? I see many arguments that Flux should be more like the other frameworks and I guess that if model building deviates the same complaints will come for that too.

[darsnack]

I'm not sure I would just blanket equate "Flux models" with "Chain". In my view, a CompGraph is also a Flux model, as is any julia function which makes use of building blocks from Flux (i.e what ONNX.jl tried to do).

This is true, and it is certainly what makes Flux so powerful. But if I need to manipulate a model, then I'd much rather operate with the same layer building blocks in Flux. There are two reasons for this:

• Flux layers are pretty well designed for manipulation with standard Julia syntax
• Most people (myself included for my work) don't want to learn another library to interpret a model

Isn't the nice thing about the Chain that if your model is super simple so that it only consists of unary operations stringed (well chained actually) you can get rid of alot of complexity as you only need to store the layers as a tuple/array since the structure is implicit from the struct itself?

I agree with this too. I don't think Flux should build graphs of the computation like other frameworks. I do think Chain should remain a simple sequentially executed list of functions. That PR doesn't necessarily invalidate that assumption. It would make representing models like Inception easier.

• Parallel is just a generic form of SkipConnection with more branches. The input and output is still a unary piece of data. In the Metalhead.jl PR that I linked, a closure is used to implement the many branches of an Inception module, since SkipConnection is limited to a two branches. If you mostly want to execute the branching module, then there is little difference here. But if you want to manipulate or access the module's parts, then having a struct/layer like `SkipConnection is useful.
• Join/Split are where it might appear like the unary structure is broken, but I would argue that the input and output are still "unary" because they are a single tuple. If you use a Split in a Chain, then the Chain only sees the one tuple as the output and not the individual pieces of the tuple. It is up to the person who used the Split in the Chain to make sure that the subsequent layer accepts the tuple as an input. We already do this when we x -> reshape(x, :) between convolution and fully-connected layers.

Generally speaking, I am not in favor of adding more layers to Flux. I think it is better to just use closures and arbitrary Julia functions whenever possible. But I think Parallel, Join, and Split have proven to be ubiquitous enough in ML that their addition is warranted. Given the assumption that the input and output to a DAG are truly unary, I think Chain + Parallel + Join + Split lets you include that DAG into a "native" Flux model. Like I mentioned, that translation is tricky and must be well-documented so people know exactly what gets interpreted as what, but the ability to make that translation would be useful for me (and I suspect many others).