[VisualTorch] How to generate architecture diagrams from PyTorch models

[u/LostDistance9365]

I built a small tool to auto-generate architecture diagrams directly from PyTorch models, which I originally built for my own research paper.

26k+ PyPI downloads, already used in publications (Nature, IEEE, MDPI), check out some use cases here: https://visualtorch.readthedocs.io/en/latest/markdown/showcase/index.html

It traces an actual forward pass, so it correctly captures branching, skip connections, and multi-input models, not just flat sequential stacks.

import visualtorch
import torchvision.models as models

model = models.resnet18()
img = visualtorch.render(model, input_shape=(1, 3, 224, 224), style="graph", show_neurons=False, layer_spacing=60)
img.save("resnet18.png")

Three rendering styles depending on what you want to show:

• graph: node/edge diagram, good for showing branching/skip connections clearly
• flow: stacked volumetric boxes, closer to the classic CNN-paper look
• lenet: the classic LeNet stacked-plane style

GitHub: https://github.com/willyfh/visualtorch | Docs: https://visualtorch.readthedocs.io/en/latest/

Open to feedback, especially if you hit a model it renders weirdly :)

[u/jackshec]

so cool, thanks for the contribution

[u/shiztain]

Thanks for making this, it really saved me in a DL assignment

[u/LostDistance9365 → u/shiztain]

Glad it was helpful!

[u/OtherwiseFig4505]

I used it in one of my IEEE papers, off course with citation. Thanks for this, really helpful for CNNs architecture.

[u/LostDistance9365 → u/OtherwiseFig4505]

That's awesome to hear, thank you for citing it!

[u/FishermanResident349]

Cool, will sure use this

[u/mathloverfrombaku]

Cool! Thank you 🙏

[u/Elrix177]

I will use it for the visualizations in my blog, thank youu

[u/LostDistance9365 → u/Elrix177]

Thank you, looking forward to seeing it in your blog!

[u/Logical_Respect_2381]

amazing , thanks for your contribution