How to convert DJL model to ONNX model?

[enpasos]

Suppose you

• set up a model in DJL and trained it with the DJL PYTORCH engine
• want to load it to a browser for inference using ONNX Runtime Web.

What would be the recommended way to convert the trained DJL model to ONNX model?

[enpasos]

To fuel the discussion I set up an example that is nicely running but far from perfect integrated with DJL.

I would appreciate advice on how to get this better integrated with DJL.

[zachgk]

In general, I'm not sure if there is a great way to do it. DJL saves models into a custom format, but only saves the parameters. The main structure of the model remains only in imperative Java code.

What it needs is some way to get the imperative code into a serialized form. This can be something more manual like how you implemented the OnnxIO interface where users can explain how a Block can be converted into serialized onnx.

Or, there are a few automatic strategies as well which you can see with the python frameworks:

1. some kind of code analysis tool which you might find under Java metaprogramming or reflection and just use it to convert the basic NDArray operations and some basic operations into their onnx equivalent
2. tracing where it runs a model but tracks what operations were performed with a graph, then converts that graph into onnx

I don't know much about Java metaprogramming, but the tracing might be doable. Some of the engines have tracing functionality built-in such as torch.jit.trace. It might be worth looking into whether DJL can use this to produce the standard pytorch serialized format. It might even be usable along with the standard torch.onnx converter as well, but regardless it should be straightforward to build a Java converter that does the same.

Another option is to build the tracing on the DJL side. We could have something like a "TracingEngine" that acts as a wrapper around another Engine but builds up the graph of called operations. Then, there just needs to be a converter from the tracing engine's call graph to onnx format. One benefit of this approach is that it should work regardless of what the underlying Engine is.

[enpasos]

Thank you for your deep insight. It helps me with my decision on how to proceed.

For the next step, I'll give it a try with my inheritance approach, where each block explicitly serializes to onnx, but reuses the serialization capability of the child blocks.

However, as you can see from my MNIST example, I introduced some ugly copy-paste blocks that I would like to eliminate.

Do you think it would be possible to open up existing DJL blocks a bit more to inheritance? This should not hurt the users of the existing block, nor would it pollute the DJL code.

See, for example, Conv2d:

If this is feasible, I would put some suggestions in a pull request.

[enpasos]

I set up suggestions on how to enhance opening for inheritance in the following pull request #1451

[frankfliu]

@enpasos

I'm not convinced that inherit from Block class is a right way to do. In the Block class, we have:

• describeInput
• getOutputShapes

I think we can improve above method to fulfill your requirement. If that doesn't work adding getters to Block should work for you.

[enpasos]

describeInput and getOutputShapes are very helpful as they provide the structural information of inputs and outputs of the blocks.

For my goal to export a model as onnx file I need two more ingredients:

• the learned (or initialized) parameters of a block
• the inner structure information of a block

The parameters are accessible by the Block interface method "ParameterList getParameters();". Perfect.

Inner structure information is for e.g.

A structure parameter like "padding" is visible within package and inheritance hierarchy, but not to the public (there are no public getters).

If you could open up the inner structure information for public read access there would be no more need for me to inherit the blocks.
This would have the advantage for a user that the export functionality could be added without any need to change djl models.

[enpasos]

One more thought: After starting to apply the onnx export pattern to a non trivial network (MuZero with a SE-ResNet based network) I learned that - from my implementation point of view - it is a straight forward implementation if every piece of the network function is a block, where each block gets the additional feature to export its function part in onnx format. Everything can be nicely encapsulated, composition is simple, and when a new function part is introduced (e.g. in a forward method or at the join function of a parallel block) the translation to onnx is in the responsibility of the block inventor. Trying to avoid to translate arbitrary functions by inspection is in my opinion much easier for the translator. That would be an argument for inheritance (or even better implementing it right away in each original block) over inspection.

[enpasos]

An example of a non-trivial network trained in DJL and exported to onnx running in your browser:
https://enpasos.ai/muzero/How#onnx