Differentiate ONNX objects

Question

I would like to make an onnx model differentiable. As I understand exporting to ONNX does not export the AutoGrad graph, is there anyway to reconstruct if after loading?

I am aware of torch-ort but to me it looks like it only works with nn.Module objects, i.e. original python pytorch models ? (see example here, here and here)

Can I in any way load an ONNX exported model and get pytorch or onnx-runtime to reconstruct the backward graph?

Alternatively, can I get onnx to export backward graph of a PyTorch nn.Module model? So that I can run it with onnx runtime?

Background: I want to work with physics based models, where I could easily write forward "energy" function, and can use their gradient ("forces") in my simulations. At present we need either numerical diff, or analytic expressions to be derived before hand.

Answer 1

You may use onnxruntime to build the backward graph. Here is some lines of code I wrote a while ago using the python API onnxruntime provides: https://github.com/sdpython/onnxcustom/blob/master/onnxcustom/training/grad_helper.py#L160.