New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support torch.fx IR visualization using netron #742
Comments
@sjain-stanford can you please share a sample file? |
@lutzroeder In the attached archive, you'll find the |
@lutzroeder Could you please check if visualizing this dot graph with netron would be possible to support? |
@sjain-stanford can you provide code samples how What is the "officially" recommended and supported approach to serialize
Can you explain why this is a problem? |
This utility is used to generate the dot graph from I did ask the PyTorch developers about serialization formats for FX Graph IR to enable Netron visualization, and here's their response:
Sure, let me explain what I mean by this:
Not all PyTorch models are exportable to ONNX or traced/scripted to TS because these only support a subset of python constructs (e.g. tracing / scripting limitations). FX symbolic trace is also tracing based, so has similar limitations but broader coverage than ONNX/TS, because it supports more of "python" constructs. At the end of the day, FX graph is used to generate pytorch model, which can go through ONNX / TS. But the IR itself can change quite a bit in this translation. For developers using FX, it would immensely help to visualize the raw FX graph without any modifications. That's where I believe this netron plugin would help. |
@sjain-stanford a better approach would be to write another script to serialize the graph into an ONNX container (or any other generic format used for deep learning that is already supported). ONNX allows for any operator type or attribute set to be stored similar to the shared |
@lutzroeder thanks for the suggestion & consideration. |
@lutzroeder I have something minimal working (FX graph >> ONNX container). For example: fx_graph.onnx.zip At the moment, the coloring metadata seems to live inside Netron and is tied to the ONNX Operator Spec. Is it possible for us to embed this coloring info as metadata in the ONNX nodes to be rendered appropriately? |
@sjain-stanford can't think of a straightforward way to do this. Some ideas:
@prasanthpul might have some recommendations. |
Thanks @lutzroeder . Assuming we could find a way to embed it in the ONNX ModelProto, I imagine it would still require changes on Netron to pull the coloring metadata for rendering. Would you be able to support that change? If so, what info should I be embedding? Any pointers would help. |
@sjain-stanford first step would be to check with ONNX folks if there is a recommended way to do this. Basically, try to standardize the ability to provide operator schemas with custom metadata into ONNX. The workaround approach that you would have to sign up to maintain, would be embedding a JSON schema like onnx-metadata.json as |
torch.fx
is a library in PyTorch 1.8 that allows python-python model transformations. It works by symbolically tracing the PyTorch model into a graph (fx.GraphModule
), which can be transformed and finally exported back to code, or used as ann.Module
directly. Currently there is no mechanism to import the graph IR into netron. An indirect path is to export to ONNX to visualize, which is not as useful if debugging transformations that potentially break ONNX exportability. It seems valuable to visualize the traced graph directly in netron.The text was updated successfully, but these errors were encountered: