Preserve original scope names in exported ONNX graph #75100
Comments
|
Hello, is anyone else assigned on this or could i have a try? |
mruberry
added
module: onnx
|
Have you tried exporting with |
I tried that first, but TL;DR it simply does not work.
class MyModel(torch.nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.chunk0 = torch.nn.Sequential(torch.nn.Conv2d(1, 1, 1), torch.nn.BatchNorm2d(1), torch.nn.ReLU())
self.chunk1 = torch.nn.Sequential(torch.nn.Conv2d(1, 1, 1), torch.nn.ReLU())
def forward(self, x):
x = self.chunk0(x)
x = self.chunk1(x)
return x
model_pt = MyModel()
input_tensor = torch.randn(1, 1, 32, 32)
torch.onnx.export(
model_pt,
input_tensor,
"model_test.onnx",
verbose=True,
input_names=["inputs"],
output_names=["outputs"],
)
model_onnx = onnx.load("model_test.onnx")
print(onnx.helper.printable_graph(model_onnx.graph))the output will look like
|
|
@marcocaccin thanks for detailed response. The The I agree the scoped names provide much better clarity. Will discuss internally at microsoft, and tracked via https://msdata.visualstudio.com/Vienna/_workitems/edit/1737678 |
|
@BowenBao thank you! |
|
Same request from @lutzroeder to make visualizing structure in Netron possible. |
@garymm, @lutzroeder Is there any example to check how |
|
This feature will be super helpful! |
|
@cchan-lm Happy to hear that you find it useful, please feel free to give us feedback! This should be ready in the next PyTorch release 1.13 which is coming soon. |
|
This is fantastic @BowenBao, thank you for your hard work and making this feature happen! |
|
@BowenBao That is awesome news, and we're looking forward to 1.13! For feedback, should we open another issue or comment here? |
Please open new issue, but feel free to mention this one for reference if needed. |
|
Currently the node names do not correspond to submodule names. For example, if I export the onnx graph for |
馃殌 The feature, motivation and pitch
When a PyTorch model is exported, the graph nodes take uninformative names like
f"{op_type}_{op_cnt}", and so do tensors (besides the input and output tensors of the overall graph).In some contexts, it would be extremely helpful if each node in the exported ONNX graph would be named after the scope name of the PyTorch module from which it came from. Additionally, intermediate tensors should also ideally be named after the node that created them.
Here's a very subjective list of why it matters:
Visualisation, Debugging
Let's say one goal of exporting to ONNX is to visualise a model (without loss of generality, let's think Netron) to reason about its architecture. For anything but the most trivial NN, this is simply not possible if nodes are named
Conv_{N}and tensors are{M}.As soon as the node name becomes
unet.encoder.stage1.conv0and its output tensorunet.encoder.stage1.conv0.output0, I become able to understand what node am I looking at, where is a skip connection going, and this sort of things.ONNX graph manipulation
For some people, exporting a model to ONNX is not the end of the story. There may be some graph- or weight- level manipulations that must be applied on sub-graphs defined by the name scopes of their nodes or their intermediate tensors: just to give an idea, a silly example here may be "apply batchnorm fusion only on the decoder of the UNet".
TensorFlow does it 馃槢
Let's be clear: I am by no means a TF fan, but credit is due here. On that side, the feature I'm pitching is the default in
tf2onnxand is a design decision that makes a lot of sense.Notes
I fully understand that name scopes may lose some of their meaning when running graph-level optimisations during the export if sets of nodes are replaced by new ones and there is no 1:1 relationship. Even in this case, though, it may be possible to assign a "common ancestor name"... or fall back to the current implementation if nothing can be done.
Alternatives
1. Fork PyTorch and diverge
Since at least some of the code of the ONNX exporting is in C++, there is no easy monkey patching solution, the library must be recompiled.
2. Hack around with a heavy hand
Someone suggested that it is possible to monkey patch the
_slow_forwardmethod oftorch.nn.Moduleto give meaningful names to the tensors flying around during graph tracing (note: not the nodes, but from output tensor names to node names the step is trivial).I tried it and it seems to be a bit of a hit and miss, I would need to have a way deeper understanding of the guts of JIT to ensure that 100% of the output tensors actually have a useful name.
Alternatively, or in conjunction, I am also trying to monkey patch
torch.onnx.utils._export()to use theinlined_graph, dump it, and rename the ONNX nodes as a post-processing.Additional context
The scope name would be readily available from the
scopeName()of eachtorch._C.Nodeof the TorchScript'sinlined_graph, but the export instead uses thegraphwhere such information has been stripped off (for example, here is the first place where the choice of using thegraphis made and becomes consequential when exporting aScriptModule).To make use of this metadata down the line, also the serialisation logic of
torch._C.Graph._export_onnx()must be adapted. Here is where the nodes of the exported ONNX proto take their names on the fly.[Edit: updated markdown header level for better rendering]
The text was updated successfully, but these errors were encountered: