After not being able to find a persistent enough serialization format I whipped together a small package which translates models into an ONNX graph:
using ONNXmutable, Flux l1 = Conv((3,3), 2=>3, relu) l2 = Dense(3, 4, elu) f = function(x,y) x = l1(x) # Home-brewed global average pool x = dropdims(mean(x, dims=(1,2)), dims=(1,2)) x = l2(x) return x + y end # Input shapes can be combinations of numbers, strings, symbols or missing x_shape = (:W, :H, 2, :Batch) y_shape = (4, :Batch) # Translate f to an ONNX graph and save it to a a file. onnx("model.onnx", f, x_shape, y_shape)
It is basically just traversing the function and creating protos and types from ONNX.jl whenever it hits a “mapped” operation.
The list of supported operations is currently nothing to brag about, but I have put some effort in making it straighforward to add mappings between julia functions and ONNX operations without modifying any code in ONNXmutable (and I’m of course more than happy to receive pull requests with more OPs).
ONNXmutable can also de-serialize ONNX models into NaiveNASflux.jl models, something which could be quite handy in a transfer learning context when one also wants to make adjustments to the model (e.g change the output size or remove/add layers).
The reason for the name of the package and the perhaps unnecessary dependency to NaiveNASflux is simply that I initially had the ambition to only do serialization/deserialization for NaiveNASflux models as I thought writing a general model translator would be too hard. However, when the dust settled I realized that the method I used works also for generic functions too as long as they are not too picky with types high up in the call hierarchy.
If the serialization functionality is useful to other people than myself but the dependency to NaiveNASflux is not I would of course be happy to move it out of this package.
So far I have not tried to load a model which was serialized using ONNXmutable in some other (non-julian) library. Please file an issue if you try this and it doesn’t work.