Graph visualization interactivity

[eric-ayllon]

I believe this project could use an interactive visualization method that allows for more advanced visualization control. For example:

• Inspecting different levels of nesting without having to visualize the deepest nesting level of all layers simultaneously.
• Showing/hiding gradients and gradient stats (size in memory, backward pass arrows, etc.)
• Save current visualization

This might not be necessary for most users, but it would surely help a lot in visualizing big and complex models without generating lots of individual visualizations of the same model.

[johnmarktaylor91]

Thanks for this feedback!

Regarding (1), there’s a vis_nesting_depth option, but it’s not super flexible at the moment (it just allows you to set a fixed number of levels). The full solution would be to allow the user to provide specific criteria for whether to nest each module. I plan to do this at some point.

For the gradient visuals, I think the full solution would be to allow the user to override some of the default text label options. TorchLens has some support for that already but I plan to make this more fully-featured.

As for saving the current visualization I’ll look into ways of doing that.

Thanks a ton for the great feedback and I hope TorchLens has been helpful for you!