Support torch.fx IR visualization using netron

[sjain-stanford]

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 a nn.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.

[lutzroeder]

@sjain-stanford can you please share a sample file?

[sjain-stanford]

@lutzroeder In the attached archive, you'll find the dot graph generated from resnet18's FX traced graph, and the corresponding svg for visualization purposes. If we could visualize this dot graph in netron, that would be super super helpful, as this opens up a new mechanism of visualizing Pytorch models directly (without going through ONNX or Torchscript). FX tracing would allow a lot more models to be visualized, without being gated by ONNX/TS. At the moment we use graphviz for rendering dot to SVG/PNG then visualize the static file. But netron is amazing. Please consider.

resnet18.zip

[sjain-stanford]

@lutzroeder Could you please check if visualizing this dot graph with netron would be possible to support?

[lutzroeder]

@sjain-stanford can you provide code samples how torch.fx models are serialized to different formats?

What is the "officially" recommended and supported approach to serialize torch.fx models?

FX tracing would allow a lot more models to be visualized, without being gated by ONNX/TS.

Can you explain why this is a problem?

[sjain-stanford]

This utility is used to generate the dot graph from torch.fx's Graph IR. See this for example usage. At the moment (because netron doesn't support visualizing this directly), they use the primitive dot visualization using pydot to svg (graphviz), which is pretty static and not as advanced as netron.

I did ask the PyTorch developers about serialization formats for FX Graph IR to enable Netron visualization, and here's their response:

fx.Graph is currently not a backwards-compatible surface (we are in beta) so we don’t want to encourage serializing those structures directly. Netron can, however, define its own serialization so long as it’s okay with potentially updating it during the next PyTorch release

Sure, let me explain what I mean by this:

FX tracing would allow a lot more models to be visualized, without being gated by ONNX/TS.

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.

[lutzroeder]

@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 .dot file.

[sjain-stanford]

@lutzroeder thanks for the suggestion & consideration.

[sjain-stanford]

@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?

[lutzroeder]

@sjain-stanford can't think of a straightforward way to do this. Some ideas:

• Add a metadata_props dictionary to OperatorProto. However, there currently isn't a way to embed OperatorSetProto into a model file.
• Workaround by storing a schema in metadata_props.

@prasanthpul might have some recommendations.

[sjain-stanford]

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.

[lutzroeder]

@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 ModelProto.metadata_props["metadata.json"]. module:name (module is the same as domain in ONNX) would be needed to identify the operator and category to define the layer category which maps to a color. Note the sample above does not define and import a domain for the newly introduced operators.