Thoughts on ONNX export

[DrChainsaw]

Since the issue about ONNX import/export turned out a bit too broad, here is an attempt to focus on the export parts.

Exporting is messy, laborious, error prone and cumbersome to test, so I think this is something one would like to be able to share, both between frameworks whenever possible (e.g. Base, LinearAlgebra and Statistics ops) and between different flavours of the same framework (e.g. canonical Flux exporter vs ONNXmutable).

One idea that comes to mind is to have a support package with only e.g. NNlib -> ONNX protos functions, but without any way to deal with the computation graph itself. That would allow for reusing the laborious mapping part irregardless of how one solves traversing the graph. It does however mean more packages to add in the global registry…

There are many more details to discuss, but I think opinions about what is described here is good for knowing what kind of packages one should create.

There are probably a lot of strategies for doing export, but here is an attempt to summarize three major strategies I can think of:

Naïve approach

Traverse data structures “from the outside”. An example of what this might look like for Flux:

function save(chain::Chain, graphproto, input_name)
    prev_name = input_name
    for layer in chain
        input_name = save(layer, graphproto, input_name)
    end
return new_unique_name(chain, graphproto)
end

function save(dense::Dense, graphproto, input_name)
	my_name = new_unique_name(dense, graphproto)
	gemm_node = NodeProto(
		inputs = [input_name, my_name *”_weigth”, my_name *”_bias”],
		outputs[my_name],
		op_type = :GEMM,
		etc…
	)
	push!(graphproto.nodes, gemm_node)
	push!(graphproto.initializer, tensorproto(dense.W, my_name * “_weight”), tensorproto(dense.b, my_name * “_bias”)) 
	activation_name = save(dense.σ, graphproto, my_name)
	return activation_name	
end
etc…

Note that this example is probably not functional, it just shows what I mean with “from the outside” above. One would probably want to pass around some Context as new things needed as input will keep appearing and it will be painful to update all the method definitions.

It has the big disadvantage that it can’t penetrate into functions. It would for example choke if it hits a SkipConnection/Parallel which concatenates outputs as there would not be a way to get the dimensions along which the concatenation shall happen.

Tracing

Traverse functions using multiple dispatch. This is what ONNXmutable does and it has the advantage over the former method that it can traverse almost any function. Here is an example from ONNXmutable:

function(l::Flux.BatchNorm)(pp::AbstractProbe)
    lname = recursename(l, nextname(pp))
    γname, βname, μname, σ²name = lname .* ("_scale", "_bias", "_mean", "_var")

    add!(pp, ONNX.NodeProto(
        input=[name(pp), γname, βname, μname, σ²name],
        output=[lname],
        name=lname,
        attribute = ONNX.AttributeProto.(["epsilon", "momentum"], [l.ϵ, l.momentum]),
        op_type="BatchNormalization"))
    add!(pp, ONNX.TensorProto(l.γ, γname))
    add!(pp, ONNX.TensorProto(l.β, βname)) # Bias not optional for batchnorm
    add!(pp, ONNX.TensorProto(l.μ, μname))
    add!(pp, ONNX.TensorProto(l.σ², σ²name))

    ppout = actfun(layertype(l), l)(newnamestrat(pp, f -> join([lname, genname(f)], "_"), lname))
    return newnamestrat(ppout, nextname(pp))
end

As you can see, its just a new method for Flux BatchNorm so that it instead of computing the result of the batchnorm op, it creates the protos and adds them to the graph.

One does not really get away from the lions share of the work, which is the manual translation from Flux struct to Protos. @ToucheSir's idea to do export in NNlib level is certainly attractive as we could cover all frameworks which make use of it.

Drawback of the above approach is that it will choke on 1) functions which are only defined for certain types and 2) control code (if,for, while).

One kind of ugly thing is that there is no room for any separate Context argument above as we can’t change the number of arguments to each function. In the example above, the AbstractProbe pp is acting as a Context, but it does look a bit goofy and non-intuitive for ops with more than one input as you generally just pick an arbitrary AbstractProbe and use it as Context.

I also sometimes see people state things like "functions shall do one single thing! If a method does something completely different for a different input type you are doing it wrong". If there is some fundamental truth to that statement then I guess this principle is in clear violation. :)

Abstract Interpretation

I don’t even know if this is the right word, and there are probably multiple levels of ways to do this. What I mean here is to work on the IR/AST whatever, something like making a compiler which compiles Julia code into ONNX ProtoBufs. This is a bit outside of my depth atm. I remember trying to just recurse and print the IRTools IR from a simple (x, y) -> cat(x,y; dims=3) function and I got something like 4 pages of output. As stated many times, Mjolnir looks like it could have been more helpful here, but it does not seem to be actively developed anymore.

This approach does not get away from the lions share of work either afaik. It just alleviates the drawbacks from the tracing approach (e.g. functions with types arguments, control code and no Context argument).