Frest start : Next steps

[ayush-1506]

Opening an issue here to keep track of this. @dfdx @DrChainsaw Would be great if you could add details about what needs to be done to add complete support for Flux here.

For context, the first set of changes were done in #46.

[DrChainsaw]

A somewhat lengthy issue here with some thoughs on how to proceed.

what needs to be done to add complete support for Flux here.

One idea is to use the package name ONNX.jl for things which are ML framework agnostic and do Flux to ONNX import/export in e.g a OnnxFlux.jl package. As discussed in the lengthy issue above, there are some other parts which can be put in separate packages (for instance NNlib <-> ONNX primitives). One obvious drawback is the overhead from many packages, including potential confusion for end users but the advantage of reuse could be beneficial when it comes to adding and testing OPs.

This issue here has a small concrete potential next step which could fit in a framework agnostic ONNX package: opus111/BaseOnnx.jl#2 but I think it needs some other eyes than mine to judge whether it is the right choice.

[ToucheSir]

I think I have a grasp on what you're proposing in that issue, but it would be great to have a couple of the motivating examples mentioned to play with :)

[dfdx]

Let's revive this issue. If I understand correctly, we are stuck on a Julia-friendly representation of ONNX graphs, either in a form of computational graph, or in a form of high-level serialization/deserialization primitives. I don't know the status of the NaiveNASlib / CompGraph, but one thing I've been working recently is the Ghost.jl package which implements the Tape type which is intentionally designed similar (though not identical) to ONNX graph (e.g. see Tape anatomy). Ghost currently supports tracing functions, tape execution, compilation to Julia code and most needed node (operation) types including loops. Long-term plan is to have ONNX interop in ~6-9 months, but whether the mapping should be direct or via another representation (e.g. CompGraph) is an open question.

I'll be happy to hear your current thoughts on the representation of ONNX graphs in Julia!

[DrChainsaw]

Sorry for being inactive here. I have had a bit of a dip w.r.t energy and motivation lately but I’m working on it :)

Ghost looks awesome and I absolutely see how I would love to throw out the ad-hoc tracing mechanism in ONNXmutable for it when it comes to exporting. Would it be fair to say it does roughly what Mjolnir would do when it comes to exporting ONNX?

How do you see the relation to import?

To me, the import requires normal (but domain specific) Julia code to just implement what the natural language and type-definitions of the ONNX spec says. For instance, to “trace” through a .onnx file, you need to do things like deserialize the strings which are the names of the inputs to a Node, look through all other nodes and initializers for those exact strings and figure out what to do with that info. Or is this something that one could/should use ghost for?

Some verbose thoughts on representation of ONNX graphs in Julia

I would say that the CompGraph in NaiveNASlib does not require any extra work w.r.t ONNX. It can be used as is, but it might not be the ideal choice in every situation (or rather, it is probably very rarely the ideal choice). I think that the Flux crowd seem to prefer that Flux import of ONNX models return a Flux.Chain, even if this means failing to import models which can’t be expressed with Chain.

Making some stand-alone CompGraph (perhaps one which is static and therefore can be typesafe) in a separate package might be useful, but it is such a trivial thing to do that it might feel silly to make a stand alone package out of it.

This is probably stating the obvious, but both CompGraph and Chain are just ways to express a function with Julia objects and thereby making it possible to write normal Julia code which reasons about and/or modifies them where a normal Julia function would require one to work with IR or some other ‘lower’ and more cumbersome representation. One reason why I think they happen to be useful for ONNX import is that the high level goal of ONNX import is to generate a function at runtime, i.e f = import(“model.onnx”) and to the extent of my knowledge (please correct me if I’m wrong) it is not possible to generate a normal Julia function like this without hitting the top level or using tools like RuntimeGeneratedFunctions.jl or GeneralizedGenerated.jl which come with a bunch of other restrictions.

All in all, I think that the choice of how to represent the model needs to be up to each individual package and I think it should be possible to do so without having to duplicate the lions share of work which is to define the ONNX primitives. Therefore it is not a central thing which needs agreement.

At least to me, there is value in being able to import an onnx-file as a model built by building blocks from ones favourite framework of choice, i.e that a Node with the operator type set to Conv in an .onnx file can be turned into a Flux.Conv or an Avalon.Conv2d and it seems obvious that this is done by separate packages.

One needs to be aware that getting full support for ONNX this way is going to be very cumbersome. For instance, a Conv operator which has both the weights and the input feature maps as trainable parameters but where the bias is computed by another node is a perfectly valid ONNX model. One could also take the approach to make a package which aims at giving full support to the ONNX spec without trying to shoehorn things into structs predefined by some other framework. This is also a valid approach but I somehow see less value in this as it would basically be a brand new framework.

The reason why I brought this up is that it does relate a bit to what assumptions one can do if one is to put some generic import code in this package. I suppose that if we do, we want to take height for the full-blown ONNX support, not just what seems to be the norm in existing frameworks even if it means that the code (and the use of it) will be more complicated.

[ToucheSir]

TL;DR copy and paste over CompGraph, hack off anything JuMP-related, add Functors.jl and let's get shipping :)

I can't speak for any of the other flux contributors, but to respond to the extended thoughts above:

I would say that the CompGraph in NaiveNASlib does not require any extra work w.r.t ONNX. It can be used as is, but it might not be the ideal choice in every situation (or rather, it is probably very rarely the ideal choice). I think that the Flux crowd seem to prefer that Flux import of ONNX models return a Flux.Chain, even if this means failing to import models which can’t be expressed with Chain.

Au contraire, model import should not rely on any framework's functionality. There's a reason both TF and PyTorch still lack any ONNX import functionality: the impedance mismatch is kind of unavoidable.

Thankfully, the Julia ecosystem means we don't have to do this! What's required to make a model that works with Flux?

1. Needs to be callable. This is handled by NNlib and its various backends.
2. Needs to be deconstructable to and reconstructable from a list of parameter arrays. This is required for functionality like eltype conversion and accelerator support. Functors.jl is a slim package that handles this, and Flux takes a dependency on it for its own layers.

ONNXMutable's CompGraph already fulfills 1), and adding support for 2) would be trivial AFAICT.

Making some stand-alone CompGraph (perhaps one which is static and therefore can be typesafe) in a separate package might be useful, but it is such a trivial thing to do that it might feel silly to make a stand alone package out of it.

Why not put it in this package? Honestly, the only thing I see preventing this from happening now is the JuMP dependency. If that and all the functionality around graph re-balancing were removed, it could be copied over pretty much wholesale.

This is probably stating the obvious, but both CompGraph and Chain are just ways to express a function with Julia objects and thereby making it possible to write normal Julia code which reasons about and/or modifies them where a normal Julia function would require one to work with IR or some other ‘lower’ and more cumbersome representation. One reason why I think they happen to be useful for ONNX import is that the high level goal of ONNX import is to generate a function at runtime, i.e f = import(“model.onnx”) and to the extent of my knowledge (please correct me if I’m wrong) it is not possible to generate a normal Julia function like this without hitting the top level or using tools like RuntimeGeneratedFunctions.jl or GeneralizedGenerated.jl which come with a bunch of other restrictions.

This is true, but the allure of runtime function generation shouldn't stop us from getting something out the door. Most users don't care about what their model looks like as long as they can import and run it. I'm not even sure that a dynamic DAG like CompGraph would be any slower than a fully generated function given how much NNlib et al. dominate computation time. Given how many questions come up on Slack/Zulip/Discourse about it, I would argue there is a bigger opportunity cost in still not having anything workable.

One needs to be aware that getting full support for ONNX this way is going to be very cumbersome. For instance, a Conv operator which has both the weights and the input feature maps as trainable parameters but where the bias is computed by another node is a perfectly valid ONNX model. One could also take the approach to make a package which aims at giving full support to the ONNX spec without trying to shoehorn things into structs predefined by some other framework. This is also a valid approach but I somehow see less value in this as it would basically be a brand new framework.

IMO full support is more important than feeling idiomatic. One of the biggest frustrations with ONNX implementations is hitting unsupported operators, because there's no good way to fix it as a user. Aesthetic concerns, on the other hand, are much easier to handle.

I also don't think such a package would be a brand new framework. Rather, it would be little more than plumbing between ONNX protobufs and thin wrappers over NNlib functions (with as much metadata as possible for higher level frameworks to use). If it turns out some op doesn't exist in NNlib, then we should add it there. Note that this and Functors.jl support would also allow frameworks like Flux to pick and choose layers to convert to more "idiomatic" representations, while still maintaining a broad level of support. ONNX.jl and CompGraph then become more of a common substrate than a separate framework.

[DrChainsaw]

Rather, it would be little more than plumbing between ONNX protobufs and thin wrappers over NNlib functions (with as much metadata as possible for higher level frameworks to use)

While I think this is a fully valid and perhaps somewhat ideal ONNX solution, I do fear that it will be quite alot of code to write and maintain. It will certainly be more code per op than the layer definitions in Flux since it has more features per op (right?). More features also mean more maintenance and looking at Flux it does not look like the layer definitions in there are low maintenance.

Some more minor comments and clarifications on CompGraph

ONNXMutable's CompGraph already fulfills 1), and adding support for 2) would be trivial AFAICT.

NaiveNASflux adds Functors.jl support to CompGraph for cpu<->gpu movement and datatype change and (barring any misconceptions I might have) it is unfortunately non-trivial and I was forced to rely on mutating the model instead of creating a new copy.

The crux is that a vertex also contains its input vertices, and the same vertex might be input to several vertices. When recursing with fmap through the graph two vertices with the same input vertex would end up with two copies of the same vertex instead after fmap.

NaiveNASlib has a copy function which kinda functions like fmap but does not have this issue, but somehow I think there has to be a reason why it wouldn't work to just implement something like fmap(f, g::CompGraph) = copy(g, f) or else I would have used it already.

Perhaps this is a reason to consider another design, e.g. one where the structure of the graph is inside the CompGraph and not inside the vertices. I went back and forth a few times on this in the prototyping stage and settled for the current solution because it makes it possible to to the size rebalancing stuff from the perspective of a single vertex.

If that and all the functionality around graph re-balancing were removed, it could be copied over pretty much wholesale.

Yeah, this would be very easy to do as the compute only parts of CompGraph are in files/structs which don't have any dependencies to JuMP or the re-balancing stuff. The LightGraphs stuff is obsolete and can just be deleted and same goes for the logging stuff.
https://github.com/DrChainsaw/NaiveNASlib.jl/blob/master/src/compgraph.jl
https://github.com/DrChainsaw/NaiveNASlib.jl/blob/master/src/vertex.jl

This is true, but the allure of runtime function generation shouldn't stop us from getting something out the door.

Fully agree. Fwiw, I have not been able to spot a significant difference between CompGraph and Chain in benchmarking. As you implied, whatever ns -> us happens due to type instability seems to disappear in the noise when the whole function is in the millisecond range. I suppose that if someone was to use e.g. NaiveGAflux for the same purpose as SymbolicRegression.jl they might find the type instability being the bottleneck but I don't lose any sleep over that.

[ToucheSir]

I would distinguish between the amount of code and the overhead of maintenance. One reason the per-layer maintenance burden is higher for Flux is because layers need to be ergonomic, robust and somewhat general. An under the hood representation needs none of those things as long as it's not exposed to users (which it shouldn't be).

As for the repeated vertex issue, I'd totally forgotten about it!. This is something we've been trying to work out recently because it affects all weight tying. See FluxML/Flux.jl#1592, FluxML/Flux.jl#1504, FluxML/Zygote.jl#991 and the last couple of days of discussion on Slack's #autodiff. I guess tracing would be a more optimal approach until we resolve weight tying, but that brings with it the issue of playing nice with source-to-source ADs like Zygote. If that's not a concern though, I say we should go for it.

Given the awkwardness of sending big paragraphs back and forth here, might I suggest @dfdx and @DrChainsaw join us this coming Tuesday for the ML Community Call? The link is at https://calendar.google.com/calendar/event?eid=X2I5N2t1ajlwNnNxMzBlOXA2b3FqMmQxbl8yMDIxMDYyMlQxNjAwMDBaIGp1bGlhbGFuZy5vcmdfa29tYXVhcWV0MTRlb2c5b2l2M3A2bzdwbWdAZw and you can copy a calendar invite into your own calendar via the links on https://julialang.org/community/.

[dfdx]

As for the repeated vertex issue, I'd totally forgotten about it!. This is something we've been trying to work out recently because it affects all weight tying. See FluxML/Flux.jl#1592, FluxML/Flux.jl#1504, FluxML/Zygote.jl#991 and the last couple of days of discussion on Slack's #autodiff.

I'm not sure this is the same issue. I think the core of the problem applying functors to CompGraph is that CompVertex contains direct (as opposed to symbolic) references to other vertices. I had this issue in early design of Ghost (back when it was part of Yota) and resolved it by having both - operations (i.e. vertex) and variables (i.e. references to operations). ONNX seems to follow the same design - its Node holds inputs as names of other nodes, not direct references. Applying a function to the whole graph then boils down to a for loop over the list of nodes instead of recursive traversal.

Given the awkwardness of sending big paragraphs back and forth here, might I suggest @dfdx and @DrChainsaw join us this coming Tuesday for the ML Community Call?

Unfortunately this is not an option for me - my open-source-working-hours are usually between 0am and 2am local time (I'm at UTC+3), and my family wouldn't really enjoy me talking aloud that late 😃 However I'll be happy to see the recording (if any).

[ToucheSir]

AIUI the main concern is with mutable structs shifting identity after fmap in a way that invalidates external references? The case of a node having multiple parents (i.e DAG vs tree structure) should be handled by fmap already, because it maintains an IdDict cache of previously visited nodes. I presume that removing the requirement for mutable vertices and any operations that would modify the graph in place (e.g. rebalancing) would make this a non-issue. Is that a fair assessment?

Unfortunately this is not an option for me...

No worries! I think it would be nice to have a touch base some time this summer, but that'll likely require some more scheduling.

[DrChainsaw]

I'm unsure as to whether I can join as it collides with dinner over here.

I suppose one good agenda item could be whether shoehorn into existing frameworks or make it fully ONNX compliant is preferable. I guess one might want to leave both doors open which I guess is doable if we proceed with the ecosystem route.

The case of a node having multiple parents (i.e DAG vs tree structure) should be handled by fmap already, because it maintains an IdDict cache of previously visited nodes.

You're right! I should probably just give it a try now and see if it works. Perhaps I could just replace the whole copy meachnism in NaiveNASlib with Functors.jl too.

FWIW, fmap works in NaiveNASflux now too (would be pretty useless if it didn't), it is just that it has the surprising side effect that it mutates the input.

[dfdx]

In a quickly evolving discussion I forgot to answer @DrChainsaw 's questions!

Ghost looks awesome and I absolutely see how I would love to throw out the ad-hoc tracing mechanism in ONNXmutable for it when it comes to exporting. Would it be fair to say it does roughly what Mjolnir would do when it comes to exporting ONNX?

A kind of, with the exception that Ghost.jl is maintained :) At some point I tried to migrate to IRTools.IR for tape representation and Mjolnir as a tracer. It didn't go well. Even with all the nice API and safeguards Julia IR (be it CodeInfo, IRTools.IR or other representations I tried) is hard to work with. Simple operations like inserting an instruction or replacing a constant with a function call may suddenly result in a segfault. Changing the tracer behavior in runtime is quite problematic (e.g. adding primitives or switching loop tracing on and off), and in general Mjolnir is designed more for partial evaluation rather than tracing. But the worst thing is the lack of a maintainer - relying on a package which is unlikely to ever see any updates is quite unreliable strategy.

To me, the import requires normal (but domain specific) Julia code to just implement what the natural language and type-definitions of the ONNX spec says.

It really depends on what your domain is. If you have a simple ONNX definition and want to construct Flux.Chain from it, then yes, direct translation sounds like a good plan. In more complex scenarios - collapsing several nodes into one or vice versa, dealing with loops, code generation, etc. - you may want an intermediate representation as a graph. Ghost provides such a graph with semantics close to that in ONNX as well as the tools for modifying this graph, execution, code generation, etc.

---

Obviously, I'm more familiar with Ghost.Tape so I keep thinking (and saying) in its terms, but I guess CompGraph should work equally well. If you have in mind the next steps which I can help with, please let me know and I'll be happy to jump in!

[ToucheSir]

Yup, I doubt Mjolnir will ever go beyond the experiment phase or see the light of day again (we should probably archive it). The work there is now being picked up and more robustly developed as part of https://github.com/JuliaCompilerPlugins and the Symbolics.jl ecosystem. I don't have a single authoritative link, but #compiler-plugins on slack has a lot of great discussion. Unlike last time, there is both upstream/core support and many downstream users (e.g. JET) involved, so it's very much worth a look to see what infrastructure can be shared.

Edit: I should add that this new effort is able to learn from the experience of IRTools/Cassette/Mjolnir development and has a much better bus factor!

[darsnack]

Flux crowd seem to prefer that Flux import of ONNX models return a Flux.Chain, even if this means failing to import models which can’t be expressed with Chain.

Potentially speaking as "the crowd" 😉, I want to echo @ToucheSir's comments and reassure that this is not quite the case. I think eventually it would be nice to have an ONNX2Flux.jl package that can translate an ONNX model to a Flux model when possible. For the simple reason that it will be in the representation that Flux.jl users expect, and it will work most easily with other packages that operate on the Flux models themselves.

But I do not think that our first step should be reading ONNX models into Flux layers. I think the approach outlined here is the right one. ONNX.jl should not rely on Flux.jl in any way. We want complete support for arbitrary ONNX structures, which means using CompGraph (or Ghost.Tape or whatever we decide) as the structure to store and run the model. If we can make it Functors.jl compatible, then that would basically ONNX2Flux.jl an aesthetic nicety in almost every case.

I am happy to help review/write code to push this effort forward. In particular, fmap should work for CompGraphs, but if you run into Functors.jl issues, then ping me/file an issue so that we can fix it.

[DrChainsaw]

Alright, I suppose the 'impedance mismatch' problem makes the separate framework the least risky way if the goal is to eventually support the whole spec although it does seem like a pretty gargantuan task to me.

Between CompGraph and Ghost, the latter seems much more worked through and generic so my vote goes for it.

Do we want the operators in this repo? I suppose they might be dead weight for e.g. ONNX2Flux.jl unless one tries to first load the model in the generic framework and then translates what can be translated from that to Flux layers.

Just some things which might be good to think about before starting:

Alot of the impedance mismatch comes from ONNX being built with row major in mind (and even when its not, models from the python frameworks will most likely be row major). Both old ONNX and ONNXNaiveNASflux (ex. ONNXmutable) tries to import as column major and the latter exports as row major. I suppose alot of risk can be mitigated if one does not attempt this.

The maintaining the versioning of the opsets might become quite burdensome if one does not think about it early (i.e before implementing a bunch of OPs). I tried asking the question in the onnx forum but did not get any replies, but this post is roughly the same and has some discussion in it: onnx/onnx#2872