Collapsing repeated blocks

[d-kleine]

Is your feature request related to a problem? Please describe.

When visualizing deep neural networks with many repeated blocks (like ResNet or Transformer layers), the generated graphs become extremely large and difficult to navigate. This is particularly problematic for models with 12+ identical layers, where the visualization becomes cluttered with redundant information.

Describe the solution you'd like

• Add an option to collapse repeated sequential blocks in the visualization with a notation indicating the number of repetitions. For example:

# Example usage
model_graph = torchview.draw_graph(
    model,
    input_size,
    collapse_repeated=True,  # New parameter
    repeat_threshold=2       # Minimum number of identical blocks to collapse
)

• The visualization would show repeated blocks as a single block with a multiplier (e.g., "Transformer Block × 12") instead of drawing each instance separately.

Describe alternatives you've considered

• Implementing a custom grouping mechanism that requires manual specification of which layers to collapse
• Adding a zoom/expand feature to hide/show repeated blocks on demand
• Creating a hierarchical view where repeated blocks are nested under a parent node

Screenshots / Text
None

Additional context

This feature would be particularly useful for:

• Transformer-based architectures (BERT, GPT, etc.)
• Deep ResNet variants
• Vision models with repeated stages
• Any architecture using nn.ModuleList with multiple identical blocks

The collapsing should be intelligent enough to detect truly identical blocks while preserving any meaningful variations in architecture.

[mert-kurttutan]

Did you check the roll argument? Does it satisfy your request ?

[d-kleine]

Thanks for your reply.

No, this applies to recursive elements. I mean repeated elements like transformer blocks that have the same components and computations inside, see for example:

from transformers import AutoModel, AutoTokenizer
from torchview import draw_graph
import torch
import graphviz
graphviz.set_jupyter_format('png')


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModel.from_pretrained("openai-community/gpt2")
tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")

inputs = tokenizer("Hello world!", return_tensors="pt")

model_graph = draw_graph(model, input_data=inputs, device=device, expand_nested=True, )

model_graph.visual_graph.node_attr["fontname"] = "Helvetica"
model_graph.visual_graph

Even though set to roll=True, every GPT2Block will be displayed individually. I think it would be great to have an option to collapse this to one block with a text next to the block's dashed line indicating the number of repetitions, like in the original transformer paper has visually the repetitions in the transformer architecture with $N\times$:

[mert-kurttutan]

I see, the example makes it clear. The distinction between the two is how we identify them, rolling is by module id (id of nn.Module object), collapsing would be by the name of module (that are contiguously applied.

For those that are interested, the starting point is to use name of nn.Module instead of id and make sure that we collapse only the ones that are contiguously applied.

We also need to make sure that collapsing and rolling are distinct visually.

Interesting features!

PRs are welcome!!

[d-kleine]

Since you know the package architecture best, I was wondering if you might be interested in implementing this feature, as described above? I believe this could be quite helpful for many users.

I'm happy to provide any additional details if that would be helpful. 🙂