Interest in adding ONNX support ?

[zstreet87]

I'm interesting in exporting a deepxde model to ONNX. This is pretty straightforward using some of the backends. For example, using PyTorch is as simple as torch.onnx.export(model, "model_saved_file.onnx",...). That's the gist of the command but needs model inputs and other arguments for protobuf (please see documentation).

https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

Since DeepXDE supports multiple backends, I figured having the functionality to transfer to one backend to another would be useful in general (for example, performance analysis with different backends). Looking at DeepXDE briefly, adding the ability to talk to ONNX on the DeepXDE layer on top of the backends doesn't seem difficult.

Any interest?

[lululxvi]

Yes, good idea. Here is the code of saving:

deepxde/deepxde/model.py

Line 815 in 532ac4c

def save(self, save_path, protocol="backend", verbose=0):

and restore:

deepxde/deepxde/model.py

Line 870 in 532ac4c

def restore(self, save_path, verbose=0):

Feel free to submit a PR.

[zstreet87]

Okay great! I've been looking at the code and I am wondering what the best way to incorporate ONNX into your design.

Each backend has a different interface to ONNX. For example, PyTorch exports an ONNX model like torch.onnx.export(model, args...) and exporting a TensorFlow model to ONNX is different. To stay in line with your design, I thought an ONNX class in models.py that is an interface to the supported ONNX backends. This would require each supported backend has the implementation code. Any thoughts and/or suggestions would be great!

[lululxvi]

I am not sure what the best way is, as I am not familiar with ONNX. What you describe seems good, i.e., we can implement the code for each backend separately.