Does Zuko allow exporting to ONNX?

[CaioDaumann]

Hello again,

Is it possible to export Zuko flows to onnx? If possible, do you have any example?

If not promptly possible, you have any ideia of much effort would it take? I would be interesting in trying that out.

Best,
Caio

[francois-rozet]

Hello @CaioDaumann,

I have never tried, but I think it should be possible. Looking at https://pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html, you will probably need to wrap the flow as a "pure function", that is something that takes tensors as input and returns tensors as output. This is not the case of the Flow objects which take a tensor as input and returns a Distribution. A very thin wrapper module that takes both $c$ and $x$ as input and returns $\log p(x | c)$ is probably enough.

Also, I don't know how ONNX handles randomness, but sampling from the flow requires first sampling from the base distribution, and then transforming it. It might be easier to only export the transform to ONNX.

In any case, if you succeed, consider contributing a tutorial to the repo!

[CaioDaumann]

Thanks for the tips @francois-rozet ! I will try it in the next days when I have some free time and will let you know if I succeed.

[dpbigler]

@CaioDaumann did you have any luck?

[CaioDaumann]

@dpbigler I haven't had time to look at this yet. I can try to do it tomorrow or Friday and let you know if I succeed.

[CaioDaumann]

Hi @francois-rozet , @dpbigler , Getting back to this, I spend some time trying to export a zuko model to onnx and I came up with something like this:

import torch
import torch.utils.data as data
import zuko
import numpy as np
import onnxruntime as ort

class WrappedNSF(torch.nn.Module):
    def __init__(self):
        super(WrappedNSF, self).__init__()
        self.flow = zuko.flows.NSF(features=2, transforms=3, hidden_features=(64, 64))
    
    def forward(self, x):
        result = self.flow().transform(x)
        return result

def two_moons(n: int, sigma: float = 1e-1):
    theta = 2 * torch.pi * torch.rand(n)
    label = (theta > torch.pi).float()

    x = torch.stack((
        torch.cos(theta) + label - 1 / 2,
        torch.sin(theta) + label / 2 - 1 / 4,
    ), axis=-1)

    return torch.normal(x, sigma), label

samples, labels = two_moons(16384)

samples_tensor = samples.clone().detach()

trainset = data.TensorDataset(samples, labels)
trainloader = data.DataLoader(trainset, batch_size=64, shuffle=True)

model = WrappedNSF()
model.eval()  

dummy_input = torch.randn(1, 2)  

output = model(dummy_input)
print("Sample output:", output)

try:
    torch.onnx.export(model,               # Wrapped model instance
                    dummy_input,         # Model input (or a tuple for multiple inputs)
                    "wrapped_flow_model.onnx",  # Output ONNX file path
                    export_params=True,    # Store the trained parameter weights inside the model file
                    opset_version=17,      # ONNX version to export the model to
                    do_constant_folding=True,  # Optimization: constant folding
                    input_names=['input'],     # Model's input names
                    output_names=['output'],   # Model's output names
                    dynamic_axes={'input': {0: 'batch_size'},  # Variable length axes
                                    'output': {0: 'batch_size'}})
    print("Model exported successfully.")
except Exception as e:
    print("Failed to export model:", str(e))

But this returns the following error:

Failed to export model: Exporting the operator 'aten::searchsorted' to ONNX opset version 17 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

Any ideias here or should I open an issue in PyTorch as the error message suggests?

[CaioDaumann]

@francois-rozet with MAF it works. Here are for example the outputs of the forward method of the flow trained in the make moons dataset. I show the output of the zuko flow, I export it to onnx and check the output of the forward method of the onnx model.

But the problem is, from what I understand onnx only supports one method, the forward one. The other ones will not be exported into the onnx model. So, every functionality you need has to be wrapped around the forward method. In order to be able to sample from the flow I did this:

class WrappedNSF(torch.nn.Module):
    def __init__(self):
        super(WrappedNSF, self).__init__()
        self.flow = zuko.flows.MAF(features=2, transforms=10, hidden_features=(64, 64))

    def forward(self, x):
        sample_flag = x[0, 0].item() > 0.5  
        num_samples = int(x[0, 1].item())  
        
        if sample_flag:
            samples = self.flow().sample((50000,))
            return samples

        x = x[:, 2:]  # Actual data starts from the third element
        result = self.flow().transform(x)
        return result

    def loss(self, x):
        return -self.flow().log_prob(x).mean()
    
    def sample(self, n):
        return self.flow().sample((n,))
    
    def get_flow(self):
        return self.flow

The input for the forward method should have a[ boolean (sample or no), n_samples, the remaining input tensor] it would have four dimensions since this is based on the make moons example. And it works well before the passing it to onnx. both the forward and sample parts of the forward method.

But for some reason, when I export it to onnx, no matter the values you add in samples = self.flow().sample((50000,)) it will always sample only one single point. I still need to understand why tho.

A quick way around it, would be to do a for loop, and sample the function n times. But I guess the seed is always the same and it always returns the same point (x and y). Is there a way to give the seed as input to the sample method?

[francois-rozet]

torch.onnx.export is based on torch.jit.trace which captures a static computation graph. This means that the computations that are performed for some input are independent from the value of that input. Therefore, torch.onnx.export will evaluate the if statements in your module once, and do the same for all inputs.

Basically, each method must have a different ONNX program. One for log_prob(x), one for transform(x), one for transform.inv(z), ...

Similarly, randomness is only evaluated once, and kept constant for all inputs. So for sampling, you should make a module that takes the latent z as input and returns flow().transform.inv(z).

[CaioDaumann]

@francois-rozet with your suggestion looks like it works nicely now. You can check the plots of the samples drawn from the onnx and the zuko model.

Regarding the searchsorted the NSF is the only type of flow that uses it?

[francois-rozet]

Nice! I believe it is the only one yes, but other flows might use other instructions that are not supported by ONNX.

[CaioDaumann]

Okay, I will test the other kind of flows to check in which types we have problems.