Package to interactively visualize large Pytorch models in notebooks - torchvista

projects
1

Hi,
I want to share a new open source package I’ve created to interactively visualize Pytorch models in notebooks (like Jupyter, Google Colab etc). It is called torchvista.

Github page

With one line of code, you can trace and visualize the forward pass of your Pytorch model.

torchvista in action

It gives you an interactive graph of the model showing how tensors are passed, and lets you explore very large models through expanding/collapsing submodules of interest.

Here are some web-based demos of real-world models if you want check them out before trying torchvista yourself: Demos page
You can also find a Google Colab tutorial on the Github page.

I hope this is useful, and am curious to get your feedback!

Thank you

3 Likes
2

How did you create the graph?

4

It had to be done by temporarily wrapping the inbuilt operations and modules to do the additional work of building the graph when they get called. Connections are made from each node to the other nodes it receives tensors from.

When the tracing is done, the operations and modules are reset back to their original definitions.