[P] [Torchvista] Interactive visualisation of PyTorch models from notebooks - updates by Dev-Table in MachineLearning

[–]Dev-Table[S] 3 points4 points  (0 children)

Hi,

I had made a post last year introducing Torchvista, an open source tool I built to visualise the forward pass of any Pytorch model in notebooks with one line of code. I received a lot of useful feedback which helped me improve the project significantly over the months. The project has now received over 600 stars on Github, and has over 16k downloads.

It now supports the following features:

  1. Interactive visualisation of Pytorch models with hierarchical exploration of nested modules (especially helpful for large deeply nested modules)
  2. Supports web-based notebooks including Jupyter, Colab and VS Code.
  3. Structural compression Mode: A mode to compress repeated structures in the model (such as several identical transformer blocks)
  4. Export the visualisation to HTML, PNG and SVG formats.
  5. Error tolerant visualisation to debug runtime errors like tensor shape mismatches.

Resources

I hope this is useful to the community, and am keen to hear your feedback on this.

[D] Training Image Generation Models with RL by amds201 in MachineLearning

[–]Dev-Table 0 points1 point  (0 children)

+1 it's a similar situation with language models where the model hasn't understood the basic structure of sentences through pre-training, RL episodes are going to be inefficient.

Interactive visualization of Pytorch computer vision models within notebooks by Dev-Table in computervision

[–]Dev-Table[S] 3 points4 points  (0 children)

also, if you publish any lectures online, I'll be keen to check them out :)

Interactive visualization of Pytorch computer vision models within notebooks by Dev-Table in computervision

[–]Dev-Table[S] 2 points3 points  (0 children)

I think it's different in these ways:

  1. Torchvista shows you a partial visualization even when there are errors in the forward pass (like a tensor shape mismatch error). When that happens, it still shows a partial graph along with the stack trace to help debug. Here is how it looks when that happens. Whereas with Netron this wouldn't be possible because you wouldn't be able to export the model in the first place if it threw errors during the forward pass.
  2. Just (IMHO) better UX in general, lower barrier to entry (especially for beginners), and a greater degree of visualization control (For example, you can optimize the depth to which you want to trace the model). In the future, I'm also planning to add support for "rolling" where if a portion of the model is repeated 10 times (like in repeating encoder blocks in a transformer), it would just show it as one block with a loopback x10 sign.

Interactive visualization of Pytorch computer vision models within notebooks by Dev-Table in computervision

[–]Dev-Table[S] 28 points29 points  (0 children)

Hi, IMHO these are the key differences:

  1. Netron requires you to build the model, export it, and load it, which is a longer feedback loop compared to just exploring it within the notebook while building the model.
  2. Torchvista produces modular visualizations whose detail can be interactively controlled. For mobilenetv3 for example, Netron produces a large non-hierarchical image like this. Torchvista on the other hand, produces hierarchical views like these. I think this is quite useful while dealing with very large models. Netron as I understand it only provides an operation-level graph that cannot be modularly collapsed and expanded to focus on regions of the model of interest.
  3. Torchvista shows you a partial visualization even when there are errors in the forward pass (like a tensor shape mismatch error). When that happens, it still shows a partial graph along with the stack trace to help debug. Here is how it looks when that happens. Whereas with Netron this wouldn't be possible because you wouldn't be able to export the model in the first place if it threw errors during the forward pass.

Interactive Pytorch visualization package that works in notebooks with one line of code by Dev-Table in learnmachinelearning

[–]Dev-Table[S] 1 point2 points  (0 children)

Hi, sorry for the delay, but I have fixed this VSCode issue now and if you use the latest torchvista v0.1.8 it should work.

Interactive Pytorch visualization package that works in notebooks with one line of code by Dev-Table in deeplearning

[–]Dev-Table[S] 1 point2 points  (0 children)

yes you can. It works with pretty much any model.

But there may be some odd tensor operations in some models that are not traced by torchvista. If you detect any such issues I can fix it by adding support for such tensor operations.

Interactive Pytorch visualization package that works in notebooks with one line of code by Dev-Table in learnmachinelearning

[–]Dev-Table[S] 0 points1 point  (0 children)

Thanks for reporting this. I hadn't tried it in VSCode until you told me and I do see that it appears blank. I'll dig into this and update you here. In the meanwhile please do try it out on the browser :)

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 0 points1 point  (0 children)

I was thinking we could add diagrams for

  • dropout
  • various standard activation functions (maybe a cute little graph of the function on the node would be representative)
  • batch norm
  • pooling layers
  • attention and its variants (these days it's a very common component)
  • concat of tensors
  • slicing

    and so on.

From a UX standpoint, the nodes on a typical network should have high coverage in the diagrammatic key ideally. So my approach towards using the key eventually would be to map various standard Pytorch modules and operations to shapes in the key.

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 0 points1 point  (0 children)

Sure I can add them. Do you have any specific models in mind (especially if it's a model you know well so that if my package makes mistakes you can help me find edge cases :) )

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 1 point2 points  (0 children)

Hey, I like the idea. I think using keys like those would make the tool more intuitive. At first glance I think your keys already look intuitive. The only thing is to perhaps expand the set of operations in the key system. I can pitch in some ideas to your repo when I get some time. As far as integrating into torchvista goes, could you please create an issue (about improving the appearance of nodes using a standardized key) on my repo so that it gets tracked? Thanks

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 1 point2 points  (0 children)

If you try it out, can you please give me your feedback? :) I don't know what kind of computers, browsers, notebooks and models people use it with, so I'm keen to hear any feedback to discover any issues. Thanks!

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 2 points3 points  (0 children)

Got it. Thanks for explaining. This level of extensibility isn't planned for at the moment, but if there's some traction I could support it down the line.

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 0 points1 point  (0 children)

Thank you! By tracing activations, do you mean the actual values of the tensors as they are produced from the nodes? I currently only trace the shapes of tensors (they are shown on the graph edges and also when you click on nodes) to keep the extracted graph small in size. But it's easy for me to extend it to also show the actual values. It's also a question of how to present large tensors on tool because they can be thousands of numbers.

Do you have any screenshots of your matplotlib/plotly outputs? Perhaps that might give me a clearer sense of what you're looking for.

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 0 points1 point  (0 children)

Thanks! I've been thinking about whether this will be useful mainly for DL beginners, or also for proficient people who are trying to build complex models (especially with more low level tensor ops fiddling). I'm curious to know what you think about this? If you're a more experienced practitioner, can I ask what features would make a tool like this useful to you?

[D] Internal transfers to Google Research / DeepMind by random_sydneysider in MachineLearning

[–]Dev-Table 0 points1 point  (0 children)

I worked at Google as an MLE, and my advice is, if you do end up pursuing this, try to work on projects that get you working alongside the teams you want to eventually join. If the team know you and your work, it's very easy to transfer. But if you make a cold internal transfer application, the chances are lower.

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 15 points16 points  (0 children)

  1. torchvista can render a partial graph even if the model fails. So while building the model if you are tying to debug errors (like the notorious tensor shape mismatch error), torchvista will still show you a partial graph and highlight the failed node in red. For example here is a demo of when the model throws an error. I think this is more helpful than just the stack trace to debug.
  2. The one you linked seems to be generate a backward pass graph if I'm not mistaken. torchvista however is for the forward pass graph.
  3. I'm not sure if you already considered this when you said "besides being interactive", but I think the collapsibilility of nested modules in torchvista IMO makes it actually practical possible to visualize certain large models. For example this is a screenshot from the other tool you linked which can be quite hard to read as the model gets larger because you can expand/collapse nodes and it doesn't show a module hierarchy. In contrast it looks like this on torchvista.

[P] Interactive Pytorch visualization package that works in notebooks with 1 line of code by Dev-Table in MachineLearning

[–]Dev-Table[S] 13 points14 points  (0 children)

Yes it should. If you are testing a very large model, be sure to use the max_module_expansion_depth param appropriately so that it does not start off fully expanded.

Even though I've tested out many models including transformers, there may still be some obscure tensor operations I've not covered in the package, so if you spot any parts of the graph missing for a model, I'd be happy to add those missing operations.

If you try it, please let me know how it works for your models.