[onnx export] toy model export to onnx

username: blackmagicianZ
cooked: <h1>
<a name="why-this-model-can-not-export-to-onnx-after-traced-some-keys-seem-missing-1" class="anchor" href="#why-this-model-can-not-export-to-onnx-after-traced-some-keys-seem-missing-1"></a>Why this model can not export to onnx? After traced some keys seem missing.</h1>
<h2>
<a name="keys-2" class="anchor" href="#keys-2"></a>Keys</h2>
<pre><code class="lang-auto">before trace
odict_keys(['weight', 'conv2d.weight', 'conv2d.bias'])

after trace
odict_keys(['weight', 'conv2d.bias'])
</code></pre>
<h2>
<a name="code-example-3" class="anchor" href="#code-example-3"></a>Code example</h2>
<pre><code class="lang-auto">import torch
from torch import nn
import torch.nn.functional as F


class toy_model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv2d = nn.Conv2d(3, 8, 3)
        self.weight = nn.Parameter(torch.rand(8, 3, 3, 3))


    def forward(self, x):
        self.weight = nn.Parameter(torch.rand(8, 6, 3, 3))
        self.conv2d.weight = self.weight
        self.conv2d.in_channels = 6

        x = self.conv2d(x)

        return x


x = torch.rand((1, 6, 64, 64))

model = toy_model()

torch.onnx.export(
    model,
    x,
    'toy.onnx',
    verbose=True,
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
)
</code></pre>
<h2>
<a name="error-4" class="anchor" href="#error-4"></a>Error</h2>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/code/toy_test.py", line 97, in &lt;module&gt;
    torch.onnx.export(
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/utils.py", line 752, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/utils.py", line 521, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/utils.py", line 465, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/sota/miniconda3/envs/sota/lib/python3.9/site-packages/torch/onnx/utils.py", line 420, in _trace_and_get_graph_from_model
    raise RuntimeError("state_dict changed after running the tracer; "
RuntimeError: state_dict changed after running the tracer; something weird is happening in your model!
</code></pre>

username: ptrblck
cooked: <p><code>onnx.export</code> will trace the model, i.e. it’ll execute the forward pass using the provided input data and will record all operations.<br>
In your <code>forward</code> method you are replacing internals of the model, which will most likely break.<br>
I don’t know if the scripting approach would work, but it might be worth a try.</p>

username: blackmagicianZ
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Did you know how to solve it or any suggestions?</p>

username: ptrblck
cooked: <p>You could try to <a href="https://pytorch.org/docs/stable/onnx.html#tracing-vs-scripting">script</a> your model, but I also don’t know if this would solve the issue or if ONNX doesn’t support model manipulation in the <code>forward</code> pass.</p>