# Exporting a model containing backpropagation to onnx

**URL:** <https://dev-discuss.pytorch.org/t/exporting-a-model-containing-backpropagation-to-onnx/1984>\
**Category:** onnx\
**Created:** [April 2, 2024, 6:15pm UTC](https://dev-discuss.pytorch.org/t/exporting-a-model-containing-backpropagation-to-onnx/1984 "2024-04-02T18:15:33Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![DavidMinorNV](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/davidminornv/32/1872_2.png) [@DavidMinorNV](https://dev-discuss.pytorch.org/u/DavidMinorNV)\
**Post date:** [April 2, 2024, 6:15pm UTC](https://dev-discuss.pytorch.org/t/exporting-a-model-containing-backpropagation-to-onnx/1984/1 "2024-04-02T18:15:33Z")

</div>

Hello,

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:

> <https://github.com/pytorch/pytorch/issues/121280>
>
> \### 🐛 Describe the bug
> 
> I'm trying to export a model to onnx involving back pr…opagation, and I've run into a number of issues. I can work around some of them, but this one seems a bit tough. You can see the kind of problem I've encountered if you run this python code:
> 
> \`\`\`python
> import torch
> 
> 
> class Model(torch.nn.Module):
> 
> def \_\_init\_\_(self):
> super().\_\_init\_\_()
> 
> self.latent\_dim = 256
> self.num\_heads = 4
> self.ff\_size=1024
> self.dropout=0.1
> self.activation="gelu"
> self.num\_layers = 4
> 
> root\_seqTransEncoderLayer = torch.nn.TransformerEncoderLayer(d\_model=self.latent\_dim,
> nhead=self.num\_heads,
> dim\_feedforward=self.ff\_size,
> dropout=self.dropout,
> activation=self.activation)
> 
> self.root\_seqTransEncoder = torch.nn.TransformerEncoder(root\_seqTransEncoderLayer,
> num\_layers=self.num\_layers)
>         
> 
> def forward(self, inputs):
> xseq = inputs\[0\]
> xseq = xseq.detach().requires\_grad\_()
> with torch.enable\_grad():
> 
> output = self.root\_seqTransEncoder(xseq)
> loss = torch.sqrt(output).sum()
> 
> return torch.autograd.grad(\[loss\], \[xseq\])\[0\]
> 
> mdl = Model()
> for p in mdl.parameters():
> p.requires\_grad\_(False)
> 
> print("export model")
> torch.onnx.export(
> Model(),
> \[torch.randn(\[20, 2, 256\]) \*\* 2\],
> "modelthing.onnx",
> input\_names=\["xseq"\],
> opset\_version=17,
> output\_names=\["lossgrad"\],
> verbose=True
> )
> \`\`\`
> 
> This probably overlaps with https://github.com/pytorch/pytorch/issues/120820 and https://github.com/pytorch/pytorch/issues/120822 as this is basically what I was trying to do when I found those bugs. There are some things I can work around (eg the backward pass of a gelu unit is unsupported but I can implement that myself using a torch.onnx api), but some things look a lot harder to work around.
> 
> This particular export fails with an error that it's trying to insert a parameter as a constant when it requires a gradient. I've turned requires\_grad off for all the model's parameters though, so I think it's erroneously trying to insert an intermediate value as a constant like it's doing in https://github.com/pytorch/pytorch/issues/120820 (I found that bug while basically trying to strip this one down). Fixing that issue will probably reveal the layer norm problem I reported here https://github.com/pytorch/pytorch/issues/120822 , the fact that the backward pass for the Gelu nonlinearity isn't implemented (at least that's something I can work around) and probably some other stuff
> 
> \### Versions
> 
> Collecting environment information...
> PyTorch version: 2.2.1+cu118
> Is debug build: False
> CUDA used to build PyTorch: 11.8
> ROCM used to build PyTorch: N/A
> 
> OS: Microsoft Windows 10 Enterprise
> GCC version: Could not collect
> Clang version: Could not collect
> CMake version: version 3.25.1
> Libc version: N/A
> 
> Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr 5 2023, 00:38:17) \[MSC v.1929 64 bit (AMD64)\] (64-bit runtime)
> Python platform: Windows-10-10.0.19045-SP0
> Is CUDA available: True
> CUDA runtime version: 10.0.130
> CUDA\_MODULE\_LOADING set to: LAZY
> GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Ti
> Nvidia driver version: 536.23
> cuDNN version: C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.8\\bin\\cudnn\_ops\_train64\_8.dll
> HIP runtime version: N/A
> MIOpen runtime version: N/A
> Is XNNPACK available: True
> 
> CPU:
> Architecture=9
> CurrentClockSpeed=3501
> DeviceID=CPU0
> Family=107
> L2CacheSize=16384
> L2CacheSpeed=
> Manufacturer=AuthenticAMD
> MaxClockSpeed=3501
> Name=AMD Ryzen Threadripper PRO 3975WX 32-Cores
> ProcessorType=3
> Revision=12544
> 
> Versions of relevant libraries:
> \[pip3\] functorch==2.0.0
> \[pip3\] lovely-numpy==0.2.8
> \[pip3\] numpy==1.24.3
> \[pip3\] onnx==1.14.1
> \[pip3\] onnx-graphsurgeon==0.3.27
> \[pip3\] onnxconverter-common==1.13.0
> \[pip3\] onnxruntime==1.15.1
> \[pip3\] optree==0.10.0
> \[pip3\] pytorch-lightning==1.4.2
> \[pip3\] tf2onnx==1.16.0
> \[pip3\] torch==2.2.1+cu118
> \[pip3\] torch-cluster==1.6.1
> \[pip3\] torch-fidelity==0.3.0
> \[pip3\] torch-geometric==2.3.0
> \[pip3\] torch-scatter==2.1.1
> \[pip3\] torch-sparse==0.6.17
> \[pip3\] torch-spline-conv==1.2.2
> \[pip3\] torchaudio==2.2.1+cu118
> \[pip3\] torchmetrics==0.6.0
> \[pip3\] torchvision==0.17.1+cu118
> \[conda\] blas 1.0 mkl
> \[conda\] cudatoolkit 11.6.0 hc0ea762\_10 conda-forge
> \[conda\] libblas 3.9.0 16\_win64\_mkl conda-forge
> \[conda\] libcblas 3.9.0 16\_win64\_mkl conda-forge
> \[conda\] liblapack 3.9.0 16\_win64\_mkl conda-forge
> \[conda\] mkl 2022.1.0 h6a75c08\_874 conda-forge
> \[conda\] mkl-include 2023.2.0 intel\_49496 intel
> \[conda\] mkl-static 2023.2.0 intel\_49496 intel
> \[conda\] numpy 1.24.2 py310hd02465a\_0 conda-forge
> \[conda\] pytorch 1.12.0 py3.10\_cuda11.6\_cudnn8\_0 pytorch
> \[conda\] pytorch-mutex 1.0 cuda pytorch

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
