How to get the node name mapping between FX and ONNX as far as possible?

siahuat0727: <p>I’m working on a project where I need to map FX node names to their corresponding ONNX node names, particularly after exporting an FX <code>GraphModule</code> to ONNX. The naive approach I’ve been using involves iterating through node types and maintaining a counter to track node names, as shown in the pseudo-code below:</p>
<pre data-code-wrap="python"><code class="lang-python">def get_onnx_to_fx_node_name_mapping(graph_module):
    counters = defaultdict(int)

    def guess_onnx_node_name(fx_node):
        onnx_node_name = FX_TO_ONNX_NODE_TYPE_MAPPING.get(str(fx_node.target))
        count = counters[onnx_node_name]
        counters[onnx_node_name] += 1
        return onnx_node_name if count == 0 else f'{onnx_node_name}_{count}'

    return {
        guess_onnx_node_name(fx_node): fx_node.name
        for fx_node in graph_module.graph.nodes
    }
</code></pre>
<p>However, certain FX nodes like <code>GELU</code> are decomposed into multiple nodes during ONNX conversion. This decomposition results in inaccurate mappings, particularly for nodes like <code>Add</code> that cannot be correctly mapped with this straightforward solution.</p>
<p>Additionally, other FX nodes like <code>aten.linear</code> may convert to either <code>Gemm</code> or <code>Matmul</code> in ONNX, depending on conditions that are currently unclear to me.</p>
<p>Here’s an example using a simple model that illustrates the issue:</p>
<pre><code class="lang-auto">import torch
from torch import nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 3, 3, 1)
        self.conv2 = nn.Conv2d(3, 3, 3, 1)
        self.gelu = nn.GELU()

    def forward(self, x):
        y1 = self.gelu(self.conv1(x))
        y2 = self.gelu(self.conv2(x))
        return y1 + y2

model = Net().eval()
example_inputs = (torch.randn(1, 3, 32, 32),)

gm = torch._export.capture_pre_autograd_graph(model, example_inputs)
gm.graph.print_tabular()

# Ignore train to enable export onnx
def train(self, *args, **kwargs):
    return self
gm.train = train.__get__(gm)
torch.onnx.export(gm, example_inputs, 'conv_gelu.onnx', opset_version=17)
</code></pre>
<p>The FX tabular output for this model is as follows:</p>
<pre><code class="lang-auto">
 opcode         name              target               args                                        kwargs
-------------  ----------------  -------------------  ------------------------------------------  --------
placeholder    arg0              arg0                 ()                                          {}
get_attr       _param_constant0  _param_constant0     ()                                          {}
get_attr       _param_constant1  _param_constant1     ()                                          {}
call_function  conv2d_default    aten.conv2d.default  (arg0, _param_constant0, _param_constant1)  {}
call_function  gelu_default      aten.gelu.default    (conv2d_default,)                           {}
get_attr       _param_constant2  _param_constant2     ()                                          {}
get_attr       _param_constant3  _param_constant3     ()                                          {}
call_function  conv2d_default_1  aten.conv2d.default  (arg0, _param_constant2, _param_constant3)  {}
call_function  gelu_default_1    aten.gelu.default    (conv2d_default_1,)                         {}
call_function  add_tensor        aten.add.Tensor      (gelu_default, gelu_default_1)              {}
output         output            output               ([add_tensor],)             
</code></pre>
<p>And the corresponding ONNX graph shows that the last <code>Add</code> node is named <code>Add_2</code> instead of <code>Add</code>, due to the decomposition of <code>GELU</code> into additional nodes. So I get a wrong mapping from <code>add_tensor</code> to <code>Add</code>.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/a/ca761fa9b13248569883cdc4976c2b1683a9722a.png" data-download-href="https://discuss.pytorch.org/uploads/default/ca761fa9b13248569883cdc4976c2b1683a9722a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/ca761fa9b13248569883cdc4976c2b1683a9722a_2_690x328.png" alt="image" data-base62-sha1="sT3wJ4uIcjFXsZRC3sltjKd1jvk" width="690" height="328" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/ca761fa9b13248569883cdc4976c2b1683a9722a_2_690x328.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/ca761fa9b13248569883cdc4976c2b1683a9722a_2_1035x492.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/ca761fa9b13248569883cdc4976c2b1683a9722a_2_1380x656.png 2x" data-dominant-color="F2F2F2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1855×883 49 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I am seeking advice on how to enhance the robustness of this mapping process. Is there an established method for accurately mapping FX node names to ONNX node names, especially in cases where nodes are decomposed or have conditional mappings?</p>