Exporting a model containing backpropagation to onnx


We’ve got a model containing a bunch of stuff like transformers, slicing, indexing using arrays, concatenation etc and, most awkwardly, a call to torch.autograd.grad(), which I’ve been trying to export to onnx for a while. I’ve made a few tickets about issues I’ve encountered along the way, many of which are linked from this one:

I was doing it via torch.onnx.export(), as there seems to be zero support for autograd stuff in torch.onnx.dynamo_export(). I’ve managed to brute force it and hack a pytorch version together so the torch.onnx.export() pathway works - I’m actually still not 100% clear if it was even meant to work in the first place…

I’d like to be able to do this with an official pytorch release though, and my hacked pytorch version isn’t really suitable for contributing to the project at the moment as 1) I don’t know if people want to do further work on torch.onnx.export() anyway 2) some of my workarounds are pretty hacky. I could probably put together a writeup of the issues I encountered and put a branch with my workarounds on github though.

There are definitely use cases for exporting models containing backprop to onnx, eg diffusion models with classifier guidance, optimizing latent codes at runtime etc, so it would be good to get proper support for it one way or another