`aot_module` + `torch.onnx.export` bug or user error? about pytorch

[thiagocrepaldi]

torch.onnx.export is based on TorchScript technilogy (aka torch.jit.script) which is not compatible with Dynamo by default

Instead, try using torch.onnx.dynamo_export to export models to ONNX based on either torch.nn.Module or Dynamo's graphs

Refer to the tutorial to get started at https://pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html

[srcarroll]

@thiagocrepaldi thanks. i did try with dynamo_export too and almost positive i ran accross the same exact issue. i'll double check. nor doees this tutorial help at all with getting it to work with autograd. i can get torch.onnx.export working fine with dynamo, but without autograd

[srcarroll]

I would be extremely grateful if you could point to an example of using aot_module with torch.onnx.dynamo_export. I've searched a lot and couldn't find one example anywhere. My assumption is that it's because it doesn't work. :P

[thiagocrepaldi]

@xadupre and @wschin can help you in more depth, as they are the maintainers for the onnx backend on torch.compile

an untested repro would be something like

import torch

class Linear(torch.nn.Module):
    def __init__(self):
        super(Linear, self).__init__()
        self.linear = torch.nn.Linear(128, 10)
        self.activation = torch.nn.ReLU()

    def forward(self, *inputs):
        input = self.linear(inputs[0])
        input = self.activation(input)
        return input

model = Linear()
model.train()
loss_fn = torch.nn.MSELoss()

input = torch.randn((64, 128), requires_grad=True)
labels = torch.randn((64, 10), requires_grad=True)

compiled_model = torch.compile(model, backend="onnxrt")
output = compiled_model(*input)
loss = loss_fn(output, labels)
loss.backward()

[thiagocrepaldi]

Refer to this https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

[srcarroll]

@thiagocrepaldi thanks again. i do know how to use torch.compile for inference, but i believe your example does not work for backward graph compilation. as i understand, that's what aot autograd is for.

it would be preferable to have something that works directly on module (either nn.Module or fx.GraphModule is fine) like aot_module is supposed to, but at this point i'd take any example that works with onnx (i have yet to see any non trivial examples). for example i've tried defining backward functions and compiling that with both torch.compile and aot_function but with no success. i look forward to hear what others suggest.