Proposal: ONNX Export Support

[corwinjoy]

Hello! We've had great success in using GPytorch for our experiments, and now we are investigating ways to streamline these models in production. We have been exporting fitted models to ONNX as part of a larger standardization effort in our company. For the most part, this has worked well, but we did have to overcome a few obstacles:
Generic Issues:

1. The floating point format is not very consistent inside GPytorch. When we went to export the computation graph, most of it used float32 calculations, but the matrix operations were largely float64 operations. We think a PR to be more explicit about numerical precision in GPytorch would make sense.
2. The library has a few places where it reshapes tensors. Most of these should not be necessary and they create difficulties for ONNX. (And probably other fixed export formats).

Export Issues:

1. The cholesky_ex operation that is used by the library is experimental. We had to create a version of _psd_safe_cholesky that just uses a regular Cholesky operation.
2. ONNX does not support linear algebra operations as part of their standard. We hope to get this added to the standard, but for now, we had to create a python script to provide the missing operations at runtime.

We would like to share our work via a PR for ONNX export from GPytorch. We think this would help others that want to export their models and may provide an impetus to make the library more precise as explained above.

Question. Is there general interest in a PR to support ONNX export?

[corwinjoy]

@dannyfriar @hughsalimbeni

[dannyfriar]

Thanks @corwinjoy!

Just a note that it's not reshaping tensors that's the problem but manipulating tensor shapes in place. This causes issues for ONNX because it traces tensor shapes as tensors.

This is the relevant section of the torch ONNX docs.

There is an example of this in-place manipulation in the block diag LO here.

[corwinjoy]

Thanks for the clarification!

…

On Thu, Jul 20, 2023, 2:08 AM Danny Friar ***@***.***> wrote: Thanks @corwinjoy <https://github.com/corwinjoy>! Just a note that it's not reshaping tensors that's the problem but manipulating tensor shapes in place. This causes issues for ONNX because it traces tensor shapes as tensors. This is the relevant section of the torch ONNX docs <https://pytorch.org/docs/stable/onnx.html#avoid-in-place-operations-when-using-tensor-shape-in-tracing-mode> . There is an example of this in-place manipulation in the block diag LO here <https://github.com/cornellius-gp/linear_operator/blob/main/linear_operator/operators/block_diag_linear_operator.py#L130> . — Reply to this email directly, view it on GitHub <#2380 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADBF42T7GHK7NZUJB4U5JUTXRDYQVANCNFSM6AAAAAA2QTWSQ4> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[Chandrachur92]

Dropping in to see if there has been any development on having a way to have ONNX exports for GPyTorch models?