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?

[mrlj-hash]

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

A vote from me at least - I see this as a major step forward for deploying GP models, or at least easing the process.

[gpleiss]

We would need some help with this task. Anyone able to help out with a PR?

[corwinjoy]

I looked into this a while ago, and one key problem is that ONNX does not have a Cholesky operator. In the past, some attempts were made to add one but they were rejected as not a core component (they are more focused on supporting deep learning graphs). So, this would have to be done along with a custom operator for Cholesky decomposition. Here is an example of a couple ways this could be implemented to help this potential PR along:
https://github.com/corwinjoy/ort_linalg/tree/chol_cpp

[jacobrgardner]

I am probably missing something here. I assume we only want to export eval to ONNX? In this case, we should not need to support Cholesky in ONNX because these are computed once and cached.

The example here for exact models and here for variational models should, possibly sans some minor issues, immediately work for export to modern ONNX with dynamo, no?

[corwinjoy]

@jacobrgardner It's been a couple years since I last looked at this. When I did this, you are mostly right that the Cholesky decomposition could be static, but sometimes it had to be recomputed due to different batch sizes. So, in general a custom operator may be required. (I think, it's been 2 years). Here is the export to ONNX that I wrote for an ApproximateGP model. Hopefully this will help in writing a PR.

https://github.com/corwinjoy/ort_linalg/blob/main/linalg_test/test_gpytorch.py